From 1af7095cb54714b514665787efd07770517df270 Mon Sep 17 00:00:00 2001 From: Concedo <39025047+LostRuins@users.noreply.github.com> Date: Sat, 21 Feb 2026 10:54:55 +0800 Subject: [PATCH] add qwen3 tts repo files --- Makefile | 2 + README.md | 1 + otherarch/qwen3tts/QWEN3TTS_LICENSE | 24 + .../qwen3tts/audio_tokenizer_decoder.cpp | 893 +++++ otherarch/qwen3tts/audio_tokenizer_decoder.h | 235 ++ .../qwen3tts/audio_tokenizer_encoder.cpp | 771 +++++ otherarch/qwen3tts/audio_tokenizer_encoder.h | 154 + otherarch/qwen3tts/coreml_code_predictor.h | 35 + .../qwen3tts/coreml_code_predictor_stub.cpp | 31 + otherarch/qwen3tts/gguf_loader.cpp | 238 ++ otherarch/qwen3tts/gguf_loader.h | 83 + otherarch/qwen3tts/main.cpp | 178 + otherarch/qwen3tts/qwen3_tts.cpp | 585 ++++ otherarch/qwen3tts/qwen3_tts.h | 150 + otherarch/qwen3tts/text_tokenizer.cpp | 351 ++ otherarch/qwen3tts/text_tokenizer.h | 88 + otherarch/qwen3tts/tts_transformer.cpp | 2899 +++++++++++++++++ otherarch/qwen3tts/tts_transformer.h | 353 ++ 18 files changed, 7071 insertions(+) create mode 100644 otherarch/qwen3tts/QWEN3TTS_LICENSE create mode 100644 otherarch/qwen3tts/audio_tokenizer_decoder.cpp create mode 100644 otherarch/qwen3tts/audio_tokenizer_decoder.h create mode 100644 otherarch/qwen3tts/audio_tokenizer_encoder.cpp create mode 100644 otherarch/qwen3tts/audio_tokenizer_encoder.h create mode 100644 otherarch/qwen3tts/coreml_code_predictor.h create mode 100644 otherarch/qwen3tts/coreml_code_predictor_stub.cpp create mode 100644 otherarch/qwen3tts/gguf_loader.cpp create mode 100644 otherarch/qwen3tts/gguf_loader.h create mode 100644 otherarch/qwen3tts/main.cpp create mode 100644 otherarch/qwen3tts/qwen3_tts.cpp create mode 100644 otherarch/qwen3tts/qwen3_tts.h create mode 100644 otherarch/qwen3tts/text_tokenizer.cpp create mode 100644 otherarch/qwen3tts/text_tokenizer.h create mode 100644 otherarch/qwen3tts/tts_transformer.cpp create mode 100644 otherarch/qwen3tts/tts_transformer.h diff --git a/Makefile b/Makefile index b3687d7e0..e2ff0bdbb 100644 --- a/Makefile +++ b/Makefile @@ -740,6 +740,8 @@ embeddingvk: examples/embedding/embedding.cpp common/arg.cpp common/speculative. $(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS) ttscppmain: otherarch/ttscpp/cli/cli.cpp otherarch/ttscpp/cli/playback.cpp otherarch/ttscpp/cli/playback.h otherarch/ttscpp/cli/write_file.cpp otherarch/ttscpp/cli/write_file.h otherarch/ttscpp/cli/vad.cpp otherarch/ttscpp/cli/vad.h otherarch/ttscpp/src/ttscpp.cpp otherarch/ttscpp/src/ttstokenizer.cpp otherarch/ttscpp/src/ttssampler.cpp otherarch/ttscpp/src/parler_model.cpp otherarch/ttscpp/src/dac_model.cpp otherarch/ttscpp/src/ttsutil.cpp otherarch/ttscpp/src/ttsargs.cpp otherarch/ttscpp/src/ttst5_encoder_model.cpp otherarch/ttscpp/src/phonemizer.cpp otherarch/ttscpp/src/tts_model.cpp otherarch/ttscpp/src/kokoro_model.cpp otherarch/ttscpp/src/dia_model.cpp otherarch/ttscpp/src/orpheus_model.cpp otherarch/ttscpp/src/snac_model.cpp otherarch/ttscpp/src/general_neural_audio_codec.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend_default.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS) $(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS) +qwen3tts: otherarch/qwen3tts/main.cpp otherarch/qwen3tts/qwen3_tts.cpp otherarch/qwen3tts/text_tokenizer.cpp otherarch/qwen3tts/gguf_loader.cpp otherarch/qwen3tts/tts_transformer.cpp otherarch/qwen3tts/audio_tokenizer_decoder.cpp otherarch/qwen3tts/audio_tokenizer_encoder.cpp otherarch/qwen3tts/coreml_code_predictor_stub.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend_default.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS) + $(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS) ggml/src/ggml-vulkan-shaders.cpp: ifdef VULKAN_BUILD diff --git a/README.md b/README.md index c9acde0c4..9a9906a39 100644 --- a/README.md +++ b/README.md @@ -184,6 +184,7 @@ and it will install everything required. Alternatively, you can download the abo - Llama.cpp source repo is at https://github.com/ggml-org/llama.cpp (MIT) - Stable-diffusion.cpp source repo is at https://github.com/leejet/stable-diffusion.cpp (MIT) - TTS.cpp source repo is at https://github.com/mmwillet/TTS.cpp (MIT) +- Qwen3TTS source repo is at https://github.com/predict-woo/qwen3-tts.cpp (MIT) - KoboldCpp source repo is at https://github.com/LostRuins/koboldcpp (AGPL) - KoboldAI Lite source repo is at https://github.com/LostRuins/lite.koboldai.net (AGPL) - For any further enquiries, contact @concedo on discord, or LostRuins on github. diff --git a/otherarch/qwen3tts/QWEN3TTS_LICENSE b/otherarch/qwen3tts/QWEN3TTS_LICENSE new file mode 100644 index 000000000..a90f38ee8 --- /dev/null +++ b/otherarch/qwen3tts/QWEN3TTS_LICENSE @@ -0,0 +1,24 @@ +The original Qwen3TTS.cpp is made by predict-woo, repo can be found at https://github.com/predict-woo/qwen3-tts.cpp +MIT License was granted at https://github.com/predict-woo/qwen3-tts.cpp/issues/4 + +MIT License + +Copyright (c) 2023-2026 The ggml authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/otherarch/qwen3tts/audio_tokenizer_decoder.cpp b/otherarch/qwen3tts/audio_tokenizer_decoder.cpp new file mode 100644 index 000000000..8129ab0cc --- /dev/null +++ b/otherarch/qwen3tts/audio_tokenizer_decoder.cpp @@ -0,0 +1,893 @@ +#include "audio_tokenizer_decoder.h" +#include "gguf_loader.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include + +#define QWEN3_TTS_DEC_MAX_NODES 32768 + +namespace qwen3_tts { + +AudioTokenizerDecoder::AudioTokenizerDecoder() = default; + +AudioTokenizerDecoder::~AudioTokenizerDecoder() { + unload_model(); +} + +void AudioTokenizerDecoder::unload_model() { + free_audio_decoder_model(model_); + + if (state_.sched) { + ggml_backend_sched_free(state_.sched); + state_.sched = nullptr; + } + if (state_.backend) { + release_preferred_backend(state_.backend); + state_.backend = nullptr; + } + if (state_.backend_cpu) { + ggml_backend_free(state_.backend_cpu); + state_.backend_cpu = nullptr; + } + + state_.compute_meta.clear(); + codes_buf_.clear(); +} + +void AudioTokenizerDecoder::normalize_codebooks() { + const float epsilon = 1e-5f; + + auto normalize_codebook = [epsilon](struct ggml_tensor * codebook, struct ggml_tensor * usage, const char *) { + if (!codebook || !usage || !codebook->data || !usage->data) return; + + int64_t codebook_dim = codebook->ne[0]; + int64_t codebook_size = codebook->ne[1]; + + ggml_fp16_t * cb_data = (ggml_fp16_t *)codebook->data; + float * usage_data = (float *)usage->data; + + for (int64_t emb_idx = 0; emb_idx < codebook_size; ++emb_idx) { + float u = usage_data[emb_idx]; + if (u < epsilon) u = epsilon; + float inv_u = 1.0f / u; + + for (int64_t dim_idx = 0; dim_idx < codebook_dim; ++dim_idx) { + int64_t mem_idx = dim_idx + emb_idx * codebook_dim; + float val = ggml_fp16_to_fp32(cb_data[mem_idx]); + cb_data[mem_idx] = ggml_fp32_to_fp16(val * inv_u); + } + } + + }; + + normalize_codebook(model_.vq_first_codebook, model_.vq_first_usage, "first"); + + for (int i = 0; i < 15; ++i) { + char name[16]; + snprintf(name, sizeof(name), "rest%d", i); + normalize_codebook(model_.vq_rest_codebook[i], model_.vq_rest_usage[i], name); + } +} + +bool AudioTokenizerDecoder::load_model(const std::string & model_path) { + unload_model(); + + GGUFLoader loader; + if (!loader.open(model_path)) { + error_msg_ = loader.get_error(); + return false; + } + + model_.config.sample_rate = loader.get_u32("qwen3-tts.tokenizer.sample_rate", 24000); + model_.config.n_codebooks = loader.get_u32("qwen3-tts.tokenizer.num_codebooks", 16); + model_.config.codebook_size = loader.get_u32("qwen3-tts.tokenizer.codebook_size", 2048); + + int64_t n_tensors = loader.get_n_tensors(); + int dec_tensor_count = 0; + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = loader.get_tensor_name(i); + if (name && strncmp(name, "tok_dec.", 8) == 0) { + dec_tensor_count++; + } + } + + if (dec_tensor_count == 0) { + error_msg_ = "No decoder tensors found in model"; + return false; + } + + size_t ctx_size = ggml_tensor_overhead() * dec_tensor_count; + struct ggml_init_params params = { + /*.mem_size =*/ ctx_size, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + model_.ctx = ggml_init(params); + if (!model_.ctx) { + error_msg_ = "Failed to initialize GGML context"; + return false; + } + + struct gguf_context * gguf_ctx = loader.get_ctx(); + struct ggml_context * meta_ctx = loader.get_meta_ctx(); + + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = loader.get_tensor_name(i); + if (!name || strncmp(name, "tok_dec.", 8) != 0) { + continue; + } + + + + struct ggml_tensor * meta_tensor = ggml_get_tensor(meta_ctx, name); + if (!meta_tensor) { + continue; + } + + struct ggml_tensor * tensor = ggml_dup_tensor(model_.ctx, meta_tensor); + ggml_set_name(tensor, name); + + model_.tensors[name] = tensor; + + std::string sname(name); + + if (sname == "tok_dec.vq_first.input_proj.weight") model_.vq_first_input_proj = tensor; + else if (sname == "tok_dec.vq_first.output_proj.weight") model_.vq_first_output_proj = tensor; + else if (sname == "tok_dec.vq_first.0.codebook") model_.vq_first_codebook = tensor; + else if (sname == "tok_dec.vq_first.0.usage") model_.vq_first_usage = tensor; + else if (sname == "tok_dec.vq_rest.input_proj.weight") model_.vq_rest_input_proj = tensor; + else if (sname == "tok_dec.vq_rest.output_proj.weight") model_.vq_rest_output_proj = tensor; + else if (sname == "tok_dec.pre_conv.weight") model_.pre_conv_w = tensor; + else if (sname == "tok_dec.pre_conv.bias") model_.pre_conv_b = tensor; + else if (sname == "tok_dec.pre_tfm.input_proj.weight") model_.pre_tfm_input_proj_w = tensor; + else if (sname == "tok_dec.pre_tfm.input_proj.bias") model_.pre_tfm_input_proj_b = tensor; + else if (sname == "tok_dec.pre_tfm.norm.weight") model_.pre_tfm_norm_w = tensor; + else if (sname == "tok_dec.pre_tfm.output_proj.weight") model_.pre_tfm_output_proj_w = tensor; + else if (sname == "tok_dec.pre_tfm.output_proj.bias") model_.pre_tfm_output_proj_b = tensor; + else if (sname == "tok_dec.dec.0.conv.weight") model_.dec0_conv_w = tensor; + else if (sname == "tok_dec.dec.0.conv.bias") model_.dec0_conv_b = tensor; + else if (sname == "tok_dec.dec.5.snake.alpha") model_.dec5_snake_alpha = tensor; + else if (sname == "tok_dec.dec.5.snake.beta") model_.dec5_snake_beta = tensor; + else if (sname == "tok_dec.dec.6.conv.weight") model_.dec6_conv_w = tensor; + else if (sname == "tok_dec.dec.6.conv.bias") model_.dec6_conv_b = tensor; + else if (sname.find("pre_tfm.blk.") != std::string::npos) { + int blk_idx; + if (sscanf(name, "tok_dec.pre_tfm.blk.%d.", &blk_idx) == 1 && blk_idx >= 0 && blk_idx < 8) { + if (sname.find(".attn_v.weight") != std::string::npos) model_.pre_tfm_layers[blk_idx].attn_v_w = tensor; + else if (sname.find(".ffn_gate.weight") != std::string::npos) model_.pre_tfm_layers[blk_idx].ffn_gate_w = tensor; + else if (sname.find(".attn_norm.weight") != std::string::npos) model_.pre_tfm_layers[blk_idx].attn_norm_w = tensor; + else if (sname.find(".attn_q.weight") != std::string::npos) model_.pre_tfm_layers[blk_idx].attn_q_w = tensor; + else if (sname.find(".attn_k.weight") != std::string::npos) model_.pre_tfm_layers[blk_idx].attn_k_w = tensor; + else if (sname.find(".attn_output.weight") != std::string::npos) model_.pre_tfm_layers[blk_idx].attn_output_w = tensor; + else if (sname.find(".attn_scale") != std::string::npos) model_.pre_tfm_layers[blk_idx].attn_scale = tensor; + else if (sname.find(".ffn_norm.weight") != std::string::npos) model_.pre_tfm_layers[blk_idx].ffn_norm_w = tensor; + else if (sname.find(".ffn_up.weight") != std::string::npos) model_.pre_tfm_layers[blk_idx].ffn_up_w = tensor; + else if (sname.find(".ffn_down.weight") != std::string::npos) model_.pre_tfm_layers[blk_idx].ffn_down_w = tensor; + else if (sname.find(".ffn_scale") != std::string::npos) model_.pre_tfm_layers[blk_idx].ffn_scale = tensor; + } + } + else { + int blk_idx, res_idx, cb_idx, n = 0; + char suffix[64]; + size_t name_len = strlen(name); + + + + #define MATCH1(fmt, var) (sscanf(name, fmt "%n", &var, &n) == 1 && (size_t)n == name_len) + #define MATCH2(fmt, v1, v2) (sscanf(name, fmt "%n", &v1, &v2, &n) == 2 && (size_t)n == name_len) + #define MATCH1S(fmt, var, suf) (sscanf(name, fmt, &var, suf) == 2) + + if (MATCH1("tok_dec.vq_rest.%d.codebook", cb_idx)) { + if (cb_idx >= 0 && cb_idx < 15) { + model_.vq_rest_codebook[cb_idx] = tensor; + } + } + else if (MATCH1("tok_dec.vq_rest.%d.usage", cb_idx)) { + if (cb_idx >= 0 && cb_idx < 15) { + model_.vq_rest_usage[cb_idx] = tensor; + } + } + else if (MATCH1S("tok_dec.upsample.%d.conv.%63s", blk_idx, suffix)) { + if (blk_idx >= 0 && blk_idx < 2) { + if (strcmp(suffix, "weight") == 0) model_.upsample[blk_idx].conv_w = tensor; + else if (strcmp(suffix, "bias") == 0) model_.upsample[blk_idx].conv_b = tensor; + } + } + else if (MATCH1S("tok_dec.upsample.%d.dwconv.%63s", blk_idx, suffix)) { + if (blk_idx >= 0 && blk_idx < 2) { + if (strcmp(suffix, "weight") == 0) model_.upsample[blk_idx].dwconv_w = tensor; + else if (strcmp(suffix, "bias") == 0) model_.upsample[blk_idx].dwconv_b = tensor; + } + } + else if (MATCH1S("tok_dec.upsample.%d.norm.%63s", blk_idx, suffix)) { + if (blk_idx >= 0 && blk_idx < 2) { + if (strcmp(suffix, "weight") == 0) model_.upsample[blk_idx].norm_w = tensor; + else if (strcmp(suffix, "bias") == 0) model_.upsample[blk_idx].norm_b = tensor; + } + } + else if (MATCH1S("tok_dec.upsample.%d.pwconv1.%63s", blk_idx, suffix)) { + if (blk_idx >= 0 && blk_idx < 2) { + if (strcmp(suffix, "weight") == 0) model_.upsample[blk_idx].pwconv1_w = tensor; + else if (strcmp(suffix, "bias") == 0) model_.upsample[blk_idx].pwconv1_b = tensor; + } + } + else if (MATCH1S("tok_dec.upsample.%d.pwconv2.%63s", blk_idx, suffix)) { + if (blk_idx >= 0 && blk_idx < 2) { + if (strcmp(suffix, "weight") == 0) model_.upsample[blk_idx].pwconv2_w = tensor; + else if (strcmp(suffix, "bias") == 0) model_.upsample[blk_idx].pwconv2_b = tensor; + } + } + else if (MATCH1("tok_dec.upsample.%d.gamma", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 2) model_.upsample[blk_idx].gamma = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.attn_norm.weight", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].attn_norm_w = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.attn_q.weight", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].attn_q_w = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.attn_k.weight", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].attn_k_w = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.attn_v.weight", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].attn_v_w = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.attn_output.weight", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].attn_output_w = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.attn_scale", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].attn_scale = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.ffn_norm.weight", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].ffn_norm_w = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.ffn_gate.weight", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].ffn_gate_w = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.ffn_up.weight", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].ffn_up_w = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.ffn_down.weight", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].ffn_down_w = tensor; + } + else if (MATCH1("tok_dec.pre_tfm.blk.%d.ffn_scale", blk_idx)) { + if (blk_idx >= 0 && blk_idx < 8) model_.pre_tfm_layers[blk_idx].ffn_scale = tensor; + } + else if (MATCH1("tok_dec.dec.%d.snake.alpha", blk_idx)) { + if (blk_idx >= 1 && blk_idx <= 4) model_.dec_blocks[blk_idx-1].snake_alpha = tensor; + } + else if (MATCH1("tok_dec.dec.%d.snake.beta", blk_idx)) { + if (blk_idx >= 1 && blk_idx <= 4) model_.dec_blocks[blk_idx-1].snake_beta = tensor; + } + else if (MATCH1("tok_dec.dec.%d.conv_t.weight", blk_idx)) { + if (blk_idx >= 1 && blk_idx <= 4) model_.dec_blocks[blk_idx-1].conv_t_w = tensor; + } + else if (MATCH1("tok_dec.dec.%d.conv_t.bias", blk_idx)) { + if (blk_idx >= 1 && blk_idx <= 4) model_.dec_blocks[blk_idx-1].conv_t_b = tensor; + } + else if (MATCH2("tok_dec.dec.%d.res.%d.act1.alpha", blk_idx, res_idx)) { + if (blk_idx >= 1 && blk_idx <= 4 && res_idx >= 2 && res_idx <= 4) { + model_.dec_blocks[blk_idx-1].res[res_idx-2].act1_alpha = tensor; + } + } + else if (MATCH2("tok_dec.dec.%d.res.%d.act1.beta", blk_idx, res_idx)) { + if (blk_idx >= 1 && blk_idx <= 4 && res_idx >= 2 && res_idx <= 4) { + model_.dec_blocks[blk_idx-1].res[res_idx-2].act1_beta = tensor; + } + } + else if (MATCH2("tok_dec.dec.%d.res.%d.conv1.weight", blk_idx, res_idx)) { + if (blk_idx >= 1 && blk_idx <= 4 && res_idx >= 2 && res_idx <= 4) { + model_.dec_blocks[blk_idx-1].res[res_idx-2].conv1_w = tensor; + } + } + else if (MATCH2("tok_dec.dec.%d.res.%d.conv1.bias", blk_idx, res_idx)) { + if (blk_idx >= 1 && blk_idx <= 4 && res_idx >= 2 && res_idx <= 4) { + model_.dec_blocks[blk_idx-1].res[res_idx-2].conv1_b = tensor; + } + } + else if (MATCH2("tok_dec.dec.%d.res.%d.act2.alpha", blk_idx, res_idx)) { + if (blk_idx >= 1 && blk_idx <= 4 && res_idx >= 2 && res_idx <= 4) { + model_.dec_blocks[blk_idx-1].res[res_idx-2].act2_alpha = tensor; + } + } + else if (MATCH2("tok_dec.dec.%d.res.%d.act2.beta", blk_idx, res_idx)) { + if (blk_idx >= 1 && blk_idx <= 4 && res_idx >= 2 && res_idx <= 4) { + model_.dec_blocks[blk_idx-1].res[res_idx-2].act2_beta = tensor; + } + } + else if (MATCH2("tok_dec.dec.%d.res.%d.conv2.weight", blk_idx, res_idx)) { + if (blk_idx >= 1 && blk_idx <= 4 && res_idx >= 2 && res_idx <= 4) { + model_.dec_blocks[blk_idx-1].res[res_idx-2].conv2_w = tensor; + } + } + else if (MATCH2("tok_dec.dec.%d.res.%d.conv2.bias", blk_idx, res_idx)) { + if (blk_idx >= 1 && blk_idx <= 4 && res_idx >= 2 && res_idx <= 4) { + model_.dec_blocks[blk_idx-1].res[res_idx-2].conv2_b = tensor; + } + } + #undef MATCH1 + #undef MATCH2 + #undef MATCH1S + } + } + + if (!load_tensor_data_from_file(model_path, gguf_ctx, model_.ctx, + model_.tensors, model_.buffer, error_msg_, + GGML_BACKEND_DEVICE_TYPE_IGPU)) { + return false; + } + + for (int i = 0; i < 4; ++i) { + model_.dec_blocks[i].res[0].dilation = 1; + model_.dec_blocks[i].res[1].dilation = 3; + model_.dec_blocks[i].res[2].dilation = 9; + } + + normalize_codebooks(); + // Codebooks are normalized in host memory; sync once to backend tensors. + auto upload_if_present = [](struct ggml_tensor * t) { + if (t && t->data) { + ggml_backend_tensor_set(t, t->data, 0, ggml_nbytes(t)); + } + }; + upload_if_present(model_.vq_first_codebook); + for (int i = 0; i < 15; ++i) { + upload_if_present(model_.vq_rest_codebook[i]); + } + + state_.backend = init_preferred_backend("AudioTokenizerDecoder", &error_msg_); + if (!state_.backend) { + return false; + } + + ggml_backend_dev_t device = ggml_backend_get_device(state_.backend); + const char * device_name = device ? ggml_backend_dev_name(device) : "Unknown"; + fprintf(stderr, " AudioTokenizerDecoder backend: %s\n", device_name); + + if (device && ggml_backend_dev_type(device) != GGML_BACKEND_DEVICE_TYPE_CPU) { + state_.backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + if (!state_.backend_cpu) { + error_msg_ = "Failed to initialize CPU fallback backend for AudioTokenizerDecoder"; + return false; + } + } + + std::vector backends; + backends.push_back(state_.backend); + if (state_.backend_cpu) { + backends.push_back(state_.backend_cpu); + } + state_.sched = ggml_backend_sched_new(backends.data(), nullptr, (int)backends.size(), QWEN3_TTS_DEC_MAX_NODES, false, true); + if (!state_.sched) { + error_msg_ = "Failed to create backend scheduler"; + return false; + } + + state_.compute_meta.resize(ggml_tensor_overhead() * QWEN3_TTS_DEC_MAX_NODES + ggml_graph_overhead()); + + return true; +} + +struct ggml_tensor * AudioTokenizerDecoder::apply_snake(struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * alpha, + struct ggml_tensor * beta) { + int64_t seq_len = x->ne[0]; + int64_t channels = x->ne[1]; + int64_t batch = x->ne[2]; + + struct ggml_tensor * alpha_exp = ggml_exp(ctx, alpha); + + struct ggml_tensor * alpha_3d = ggml_reshape_3d(ctx, alpha_exp, 1, channels, 1); + struct ggml_tensor * alpha_broad = ggml_repeat(ctx, alpha_3d, + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, seq_len, channels, batch)); + + struct ggml_tensor * ax = ggml_mul(ctx, x, alpha_broad); + struct ggml_tensor * sin_ax = ggml_sin(ctx, ax); + struct ggml_tensor * sin_sq = ggml_sqr(ctx, sin_ax); + + struct ggml_tensor * neg_beta = ggml_scale(ctx, beta, -1.0f); + struct ggml_tensor * inv_beta_exp = ggml_exp(ctx, neg_beta); + struct ggml_tensor * inv_beta_3d = ggml_reshape_3d(ctx, inv_beta_exp, 1, channels, 1); + struct ggml_tensor * inv_beta = ggml_repeat(ctx, inv_beta_3d, + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, seq_len, channels, batch)); + + struct ggml_tensor * scaled_sin = ggml_mul(ctx, sin_sq, inv_beta); + + return ggml_add(ctx, x, scaled_sin); +} + +struct ggml_tensor * AudioTokenizerDecoder::apply_rms_norm(struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * w, + float eps) { + struct ggml_tensor * normed = ggml_rms_norm(ctx, x, eps); + return ggml_mul(ctx, normed, w); +} + +struct ggml_tensor * AudioTokenizerDecoder::apply_pre_tfm_layer(struct ggml_context * ctx, + struct ggml_tensor * x, + const pre_tfm_layer & layer, + int32_t n_frames, + struct ggml_tensor * positions) { + const auto & cfg = model_.config; + const int n_heads = cfg.n_heads; + const int qkv_dim = cfg.latent_dim; + const int head_dim = qkv_dim / n_heads; + + if (!layer.attn_norm_w || !layer.attn_q_w || !layer.attn_k_w || !layer.attn_v_w || + !layer.attn_output_w || !layer.ffn_norm_w || !layer.ffn_gate_w || + !layer.ffn_up_w || !layer.ffn_down_w) { + return x; + } + + struct ggml_tensor * residual = x; + + struct ggml_tensor * normed = apply_rms_norm(ctx, x, layer.attn_norm_w, cfg.rms_norm_eps); + + struct ggml_tensor * Qcur = ggml_mul_mat(ctx, layer.attn_q_w, normed); + struct ggml_tensor * Kcur = ggml_mul_mat(ctx, layer.attn_k_w, normed); + struct ggml_tensor * Vcur = ggml_mul_mat(ctx, layer.attn_v_w, normed); + + Qcur = ggml_reshape_3d(ctx, Qcur, head_dim, n_heads, n_frames); + Kcur = ggml_reshape_3d(ctx, Kcur, head_dim, n_heads, n_frames); + Vcur = ggml_reshape_3d(ctx, Vcur, head_dim, n_heads, n_frames); + + Qcur = ggml_rope_ext(ctx, Qcur, positions, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + cfg.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + Kcur = ggml_rope_ext(ctx, Kcur, positions, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + cfg.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + struct ggml_tensor * Q = ggml_permute(ctx, Qcur, 0, 2, 1, 3); + struct ggml_tensor * K = ggml_permute(ctx, Kcur, 0, 2, 1, 3); + struct ggml_tensor * V = ggml_permute(ctx, Vcur, 0, 2, 1, 3); + + struct ggml_tensor * KQ = ggml_mul_mat(ctx, K, Q); + KQ = ggml_scale(ctx, KQ, 1.0f / sqrtf((float)head_dim)); + // Apply causal mask (each position can only attend to itself and previous positions) + KQ = ggml_diag_mask_inf(ctx, KQ, 0); + KQ = ggml_soft_max(ctx, KQ); + + V = ggml_cont(ctx, ggml_transpose(ctx, V)); + + struct ggml_tensor * KQV = ggml_mul_mat(ctx, V, KQ); + KQV = ggml_permute(ctx, KQV, 0, 2, 1, 3); + struct ggml_tensor * attn_out = ggml_cont_2d(ctx, KQV, n_heads * head_dim, n_frames); + + attn_out = ggml_mul_mat(ctx, layer.attn_output_w, attn_out); + + if (layer.attn_scale) { + attn_out = ggml_mul(ctx, attn_out, layer.attn_scale); + } + + x = ggml_add(ctx, residual, attn_out); + residual = x; + + normed = apply_rms_norm(ctx, x, layer.ffn_norm_w, cfg.rms_norm_eps); + + struct ggml_tensor * gate = ggml_mul_mat(ctx, layer.ffn_gate_w, normed); + struct ggml_tensor * up = ggml_mul_mat(ctx, layer.ffn_up_w, normed); + + gate = ggml_silu(ctx, gate); + struct ggml_tensor * ffn_out = ggml_mul(ctx, gate, up); + + ffn_out = ggml_mul_mat(ctx, layer.ffn_down_w, ffn_out); + + if (layer.ffn_scale) { + ffn_out = ggml_mul(ctx, ffn_out, layer.ffn_scale); + } + + return ggml_add(ctx, residual, ffn_out); +} + +struct ggml_tensor * AudioTokenizerDecoder::apply_upsample_block(struct ggml_context * ctx, + struct ggml_tensor * x, + const upsample_block & block, + int block_idx) { + int64_t seq_len = x->ne[0]; + int64_t channels = x->ne[1]; + + struct ggml_tensor * x_2d = ggml_reshape_2d(ctx, x, seq_len, channels); + x_2d = ggml_conv_transpose_1d(ctx, block.conv_w, x_2d, 2, 0, 1); + + int64_t new_seq_len = x_2d->ne[0]; + x = ggml_reshape_3d(ctx, x_2d, new_seq_len, channels, 1); + + if (block.conv_b) { + x = ggml_add(ctx, x, ggml_reshape_3d(ctx, block.conv_b, 1, channels, 1)); + } + + struct ggml_tensor * residual = x; + + if (block.dwconv_w) { + // Causal padding: pad left with 6 zeros (kernel_size - 1 = 7 - 1 = 6) + x = ggml_pad_ext(ctx, x, 6, 0, 0, 0, 0, 0, 0, 0); // left pad only + x = ggml_conv_1d_dw(ctx, block.dwconv_w, x, 1, 0, 1); // no padding in conv + if (block.dwconv_b) { + x = ggml_add(ctx, x, ggml_reshape_3d(ctx, block.dwconv_b, 1, channels, 1)); + } + } + + x = ggml_permute(ctx, x, 1, 0, 2, 3); + x = ggml_cont(ctx, x); + + if (block.norm_w && block.norm_b) { + x = ggml_norm(ctx, x, 1e-6f); + x = ggml_mul(ctx, x, block.norm_w); + x = ggml_add(ctx, x, block.norm_b); + } + + x = ggml_mul_mat(ctx, block.pwconv1_w, x); + if (block.pwconv1_b) { + x = ggml_add(ctx, x, block.pwconv1_b); + } + + x = ggml_gelu(ctx, x); + + x = ggml_mul_mat(ctx, block.pwconv2_w, x); + if (block.pwconv2_b) { + x = ggml_add(ctx, x, block.pwconv2_b); + } + + x = ggml_permute(ctx, x, 1, 0, 2, 3); + x = ggml_cont(ctx, x); + + if (block.gamma) { + struct ggml_tensor * gamma_3d = ggml_reshape_3d(ctx, block.gamma, 1, channels, 1); + x = ggml_mul(ctx, x, ggml_repeat(ctx, gamma_3d, + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, new_seq_len, channels, 1))); + } + + return ggml_add(ctx, residual, x); +} + +struct ggml_tensor * AudioTokenizerDecoder::apply_residual_block(struct ggml_context * ctx, + struct ggml_tensor * x, + const residual_block & block) { + struct ggml_tensor * residual = x; + + if (block.act1_alpha) { + x = apply_snake(ctx, x, block.act1_alpha, block.act1_beta); + } + + int64_t out_channels = block.conv1_w->ne[2]; + int padding = 6 * block.dilation; + x = ggml_pad_ext(ctx, x, padding, 0, 0, 0, 0, 0, 0, 0); + x = ggml_conv_1d(ctx, block.conv1_w, x, 1, 0, block.dilation); + if (block.conv1_b) { + x = ggml_add(ctx, x, ggml_reshape_3d(ctx, block.conv1_b, 1, out_channels, 1)); + } + + if (block.act2_alpha) { + x = apply_snake(ctx, x, block.act2_alpha, block.act2_beta); + } + + out_channels = block.conv2_w->ne[2]; + x = ggml_conv_1d(ctx, block.conv2_w, x, 1, 0, 1); + if (block.conv2_b) { + x = ggml_add(ctx, x, ggml_reshape_3d(ctx, block.conv2_b, 1, out_channels, 1)); + } + + return ggml_add(ctx, residual, x); +} + +struct ggml_tensor * AudioTokenizerDecoder::apply_decoder_block(struct ggml_context * ctx, + struct ggml_tensor * x, + const decoder_block & block, + int upsample_rate, + int block_idx) { + if (block.snake_alpha && block.snake_beta) { + x = apply_snake(ctx, x, block.snake_alpha, block.snake_beta); + } + + int64_t seq_len = x->ne[0]; + int64_t in_channels = x->ne[1]; + int64_t out_channels = block.conv_t_w->ne[1]; + int kernel_size = block.conv_t_w->ne[0]; + + struct ggml_tensor * x_2d = ggml_reshape_2d(ctx, x, seq_len, in_channels); + x_2d = ggml_conv_transpose_1d(ctx, block.conv_t_w, x_2d, upsample_rate, 0, 1); + + int64_t new_seq_len = x_2d->ne[0]; + x = ggml_reshape_3d(ctx, x_2d, new_seq_len, out_channels, 1); + + // Python CausalTransConvNet: left_pad = right_pad = kernel_size - stride + int pad = kernel_size - upsample_rate; + int left_pad = pad; + int right_pad = pad; + int64_t out_seq_len = new_seq_len - left_pad - right_pad; + + x = ggml_view_3d(ctx, x, out_seq_len, out_channels, 1, + x->nb[1], x->nb[2], left_pad * x->nb[0]); + x = ggml_cont(ctx, x); + + if (block.conv_t_b) { + x = ggml_add(ctx, x, ggml_reshape_3d(ctx, block.conv_t_b, 1, out_channels, 1)); + } + + for (int i = 0; i < 3; ++i) { + x = apply_residual_block(ctx, x, block.res[i]); + } + + return x; +} + +struct ggml_cgraph * AudioTokenizerDecoder::build_graph(int32_t n_frames) { + const auto & cfg = model_.config; + + struct ggml_init_params params = { + /*.mem_size =*/ state_.compute_meta.size(), + /*.mem_buffer =*/ state_.compute_meta.data(), + /*.no_alloc =*/ true, + }; + + struct ggml_context * ctx0 = ggml_init(params); + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_DEC_MAX_NODES, false); + + static const char * cb_names[16] = { + "codes_cb0", "codes_cb1", "codes_cb2", "codes_cb3", + "codes_cb4", "codes_cb5", "codes_cb6", "codes_cb7", + "codes_cb8", "codes_cb9", "codes_cb10", "codes_cb11", + "codes_cb12", "codes_cb13", "codes_cb14", "codes_cb15" + }; + + struct ggml_tensor * cb_codes_tensors[16]; + for (int cb = 0; cb < 16; ++cb) { + cb_codes_tensors[cb] = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_frames); + ggml_set_name(cb_codes_tensors[cb], cb_names[cb]); + ggml_set_input(cb_codes_tensors[cb]); + } + + struct ggml_tensor * first_codes = cb_codes_tensors[0]; + + struct ggml_tensor * first_emb = ggml_get_rows(ctx0, model_.vq_first_codebook, first_codes); + ggml_set_name(first_emb, "first_emb_raw"); + + struct ggml_tensor * rest_emb[15]; + for (int cb = 0; cb < 15; ++cb) { + struct ggml_tensor * cb_codes = cb_codes_tensors[cb + 1]; + rest_emb[cb] = ggml_get_rows(ctx0, model_.vq_rest_codebook[cb], cb_codes); + + if (cb == 0) { + ggml_set_name(rest_emb[cb], "rest_cb0_emb_raw"); + } + } + + struct ggml_tensor * first_emb_2d = ggml_reshape_2d(ctx0, first_emb, cfg.codebook_dim, n_frames); + ggml_set_name(first_emb_2d, "first_emb_2d"); + + struct ggml_tensor * first_proj_weight_2d = ggml_reshape_2d(ctx0, model_.vq_first_output_proj, + cfg.codebook_dim, cfg.hidden_dim); + struct ggml_tensor * first_proj_2d = ggml_mul_mat(ctx0, first_proj_weight_2d, first_emb_2d); + ggml_set_name(first_proj_2d, "first_proj_2d"); + + struct ggml_tensor * rest_proj_weight_2d = ggml_reshape_2d(ctx0, model_.vq_rest_output_proj, + cfg.codebook_dim, cfg.hidden_dim); + + struct ggml_tensor * rest_proj_2d = nullptr; + for (int cb = 0; cb < 15; ++cb) { + struct ggml_tensor * cb_emb_2d = ggml_reshape_2d(ctx0, rest_emb[cb], cfg.codebook_dim, n_frames); + + if (cb == 0) { + ggml_set_name(cb_emb_2d, "rest_cb0_emb_2d"); + } + + struct ggml_tensor * cb_proj_2d = ggml_mul_mat(ctx0, rest_proj_weight_2d, cb_emb_2d); + + if (rest_proj_2d == nullptr) { + rest_proj_2d = cb_proj_2d; + } else { + rest_proj_2d = ggml_add(ctx0, rest_proj_2d, cb_proj_2d); + } + } + ggml_set_name(rest_proj_2d, "rest_proj_2d"); + + struct ggml_tensor * latent_2d = ggml_add(ctx0, first_proj_2d, rest_proj_2d); + ggml_set_name(latent_2d, "latent_2d"); + + struct ggml_tensor * latent_t = ggml_transpose(ctx0, latent_2d); + ggml_set_name(latent_t, "latent_t"); + + struct ggml_tensor * latent_cont = ggml_cont(ctx0, latent_t); + ggml_set_name(latent_cont, "latent_cont"); + + struct ggml_tensor * latent = ggml_reshape_3d(ctx0, latent_cont, n_frames, cfg.hidden_dim, 1); + + ggml_set_name(latent, "vq_output"); + + struct ggml_tensor * latent_for_conv = ggml_cont(ctx0, latent); + struct ggml_tensor * latent_padded = ggml_pad_ext(ctx0, latent_for_conv, 2, 0, 0, 0, 0, 0, 0, 0); + struct ggml_tensor * cur = ggml_conv_1d(ctx0, model_.pre_conv_w, latent_padded, 1, 0, 1); + if (model_.pre_conv_b) { + cur = ggml_add(ctx0, cur, ggml_reshape_3d(ctx0, model_.pre_conv_b, 1, cfg.latent_dim, 1)); + } + + ggml_set_name(cur, "pre_conv_output"); + + struct ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur, n_frames, cfg.latent_dim); + struct ggml_tensor * cur_t = ggml_transpose(ctx0, cur_2d); + cur = ggml_cont(ctx0, cur_t); + + ggml_set_name(cur, "pre_conv_reshaped"); + + cur = ggml_mul_mat(ctx0, model_.pre_tfm_input_proj_w, cur); + if (model_.pre_tfm_input_proj_b) { + cur = ggml_add(ctx0, cur, model_.pre_tfm_input_proj_b); + } + + ggml_set_name(cur, "pre_tfm_input"); + + struct ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_frames); + ggml_set_name(positions, "positions"); + ggml_set_input(positions); + + for (int i = 0; i < cfg.n_pre_tfm_layers; ++i) { + cur = apply_pre_tfm_layer(ctx0, cur, model_.pre_tfm_layers[i], n_frames, positions); + } + + if (model_.pre_tfm_norm_w) { + cur = apply_rms_norm(ctx0, cur, model_.pre_tfm_norm_w, cfg.rms_norm_eps); + } + + cur = ggml_mul_mat(ctx0, model_.pre_tfm_output_proj_w, cur); + if (model_.pre_tfm_output_proj_b) { + cur = ggml_add(ctx0, cur, model_.pre_tfm_output_proj_b); + } + + ggml_set_name(cur, "pre_tfm_output"); + + cur = ggml_permute(ctx0, cur, 1, 0, 2, 3); + cur = ggml_cont(ctx0, cur); + cur = ggml_reshape_3d(ctx0, cur, n_frames, cfg.latent_dim, 1); + + ggml_set_name(cur, "pre_tfm_reshaped"); + + for (int i = 0; i < 2; ++i) { + cur = apply_upsample_block(ctx0, cur, model_.upsample[i], i); + } + + ggml_set_name(cur, "upsample_output"); + + // Causal padding: left pad with 6 (kernel_size - 1 = 7 - 1 = 6) + cur = ggml_pad_ext(ctx0, cur, 6, 0, 0, 0, 0, 0, 0, 0); + cur = ggml_conv_1d(ctx0, model_.dec0_conv_w, cur, 1, 0, 1); + if (model_.dec0_conv_b) { + cur = ggml_add(ctx0, cur, ggml_reshape_3d(ctx0, model_.dec0_conv_b, 1, cfg.decoder_dim, 1)); + } + + ggml_set_name(cur, "dec0_output"); + + int upsample_rates[4] = {8, 5, 4, 3}; + for (int i = 0; i < 4; ++i) { + cur = apply_decoder_block(ctx0, cur, model_.dec_blocks[i], upsample_rates[i], i); + char name[32]; + snprintf(name, sizeof(name), "dec%d_output", i + 1); + ggml_set_name(cur, name); + } + + if (model_.dec5_snake_alpha) { + cur = apply_snake(ctx0, cur, model_.dec5_snake_alpha, model_.dec5_snake_beta); + } + + ggml_set_name(cur, "dec5_output"); + + // Causal padding: left pad with 6 (kernel_size - 1 = 7 - 1 = 6) + cur = ggml_pad_ext(ctx0, cur, 6, 0, 0, 0, 0, 0, 0, 0); + cur = ggml_conv_1d(ctx0, model_.dec6_conv_w, cur, 1, 0, 1); + if (model_.dec6_conv_b) { + cur = ggml_add(ctx0, cur, ggml_reshape_3d(ctx0, model_.dec6_conv_b, 1, 1, 1)); + } + + ggml_set_name(cur, "dec6_output"); + + cur = ggml_tanh(ctx0, cur); + + cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]); + + ggml_set_name(cur, "audio"); + ggml_set_output(cur); + + ggml_build_forward_expand(gf, cur); + + ggml_free(ctx0); + + return gf; +} + +bool AudioTokenizerDecoder::decode(const int32_t * codes, int32_t n_frames, + std::vector & samples) { + if (!model_.ctx) { + error_msg_ = "Model not loaded"; + return false; + } + + const auto & cfg = model_.config; + + codes_buf_.resize(n_frames * cfg.n_codebooks); + for (int f = 0; f < n_frames; ++f) { + for (int cb = 0; cb < cfg.n_codebooks; ++cb) { + codes_buf_[cb + f * cfg.n_codebooks] = codes[f * cfg.n_codebooks + cb]; + } + } + + struct ggml_cgraph * gf = build_graph(n_frames); + + if (!ggml_backend_sched_alloc_graph(state_.sched, gf)) { + error_msg_ = "Failed to allocate graph"; + return false; + } + + std::vector cb_codes(n_frames); + for (int cb = 0; cb < 16; ++cb) { + char name[32]; + snprintf(name, sizeof(name), "codes_cb%d", cb); + struct ggml_tensor * cb_tensor = ggml_graph_get_tensor(gf, name); + if (!cb_tensor) { + error_msg_ = "Failed to find codes tensor for codebook " + std::to_string(cb); + ggml_backend_sched_reset(state_.sched); + return false; + } + + for (int f = 0; f < n_frames; ++f) { + cb_codes[f] = codes_buf_[f * cfg.n_codebooks + cb]; + } + + ggml_backend_tensor_set(cb_tensor, cb_codes.data(), 0, n_frames * sizeof(int32_t)); + } + + + + struct ggml_tensor * positions_tensor = ggml_graph_get_tensor(gf, "positions"); + if (positions_tensor) { + std::vector positions(n_frames); + for (int i = 0; i < n_frames; ++i) { + positions[i] = i; + } + ggml_backend_tensor_set(positions_tensor, positions.data(), 0, + n_frames * sizeof(int32_t)); + } + + + + if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) { + error_msg_ = "Failed to compute graph"; + ggml_backend_sched_reset(state_.sched); + return false; + } + + struct ggml_tensor * audio_tensor = ggml_graph_get_tensor(gf, "audio"); + if (!audio_tensor) { + error_msg_ = "Failed to find audio tensor"; + ggml_backend_sched_reset(state_.sched); + return false; + } + + int64_t n_samples = audio_tensor->ne[0]; + samples.resize(n_samples); + ggml_backend_tensor_get(audio_tensor, samples.data(), 0, n_samples * sizeof(float)); + + ggml_backend_sched_reset(state_.sched); + + return true; +} + +void free_audio_decoder_model(audio_decoder_model & model) { + if (model.buffer) { + ggml_backend_buffer_free(model.buffer); + model.buffer = nullptr; + } + if (model.ctx) { + ggml_free(model.ctx); + model.ctx = nullptr; + } + model.tensors.clear(); +} + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/audio_tokenizer_decoder.h b/otherarch/qwen3tts/audio_tokenizer_decoder.h new file mode 100644 index 000000000..03d734416 --- /dev/null +++ b/otherarch/qwen3tts/audio_tokenizer_decoder.h @@ -0,0 +1,235 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" +#include "gguf.h" + +#include +#include +#include +#include + +namespace qwen3_tts { + +// Audio tokenizer decoder (vocoder) configuration +struct audio_decoder_config { + int32_t sample_rate = 24000; + int32_t n_codebooks = 16; // Total codebooks (1 first + 15 rest) + int32_t codebook_size = 2048; // Entries per codebook + int32_t codebook_dim = 256; // Embedding dimension per codebook + int32_t latent_dim = 1024; // Latent dimension after VQ + int32_t hidden_dim = 512; // Pre-transformer hidden dimension + int32_t n_pre_tfm_layers = 8; // Pre-transformer layers + int32_t n_heads = 16; // Attention heads in pre-transformer + int32_t ffn_dim = 1024; // FFN intermediate dimension + int32_t decoder_dim = 1536; // Initial decoder dimension + int32_t upsample_rates[4] = {8, 5, 4, 3}; // Total: 480x upsampling + float rms_norm_eps = 1e-5f; + float rope_theta = 10000.0f; +}; + +// Pre-transformer layer weights +struct pre_tfm_layer { + // Attention + struct ggml_tensor * attn_norm_w = nullptr; + struct ggml_tensor * attn_q_w = nullptr; + struct ggml_tensor * attn_k_w = nullptr; + struct ggml_tensor * attn_v_w = nullptr; + struct ggml_tensor * attn_output_w = nullptr; + struct ggml_tensor * attn_scale = nullptr; // layer_scale for attention + + // FFN (SwiGLU) + struct ggml_tensor * ffn_norm_w = nullptr; + struct ggml_tensor * ffn_gate_w = nullptr; + struct ggml_tensor * ffn_up_w = nullptr; + struct ggml_tensor * ffn_down_w = nullptr; + struct ggml_tensor * ffn_scale = nullptr; // layer_scale for FFN +}; + +// Residual block weights (Snake + Conv + Snake + Conv) +struct residual_block { + int dilation = 1; // Dilation for conv1: [1, 3, 9] for res[0], res[1], res[2] + struct ggml_tensor * act1_alpha = nullptr; + struct ggml_tensor * act1_beta = nullptr; + struct ggml_tensor * conv1_w = nullptr; + struct ggml_tensor * conv1_b = nullptr; + struct ggml_tensor * act2_alpha = nullptr; + struct ggml_tensor * act2_beta = nullptr; + struct ggml_tensor * conv2_w = nullptr; + struct ggml_tensor * conv2_b = nullptr; +}; + +// Decoder block weights (Snake + ConvTranspose + Residual blocks) +struct decoder_block { + // Snake activation before conv transpose + struct ggml_tensor * snake_alpha = nullptr; + struct ggml_tensor * snake_beta = nullptr; + + // Transposed convolution for upsampling + struct ggml_tensor * conv_t_w = nullptr; + struct ggml_tensor * conv_t_b = nullptr; + + // Residual blocks (3 per decoder block) + residual_block res[3]; +}; + +// Upsample block weights (ConvNeXt-style) +struct upsample_block { + struct ggml_tensor * conv_w = nullptr; + struct ggml_tensor * conv_b = nullptr; + struct ggml_tensor * dwconv_w = nullptr; + struct ggml_tensor * dwconv_b = nullptr; + struct ggml_tensor * norm_w = nullptr; + struct ggml_tensor * norm_b = nullptr; + struct ggml_tensor * pwconv1_w = nullptr; + struct ggml_tensor * pwconv1_b = nullptr; + struct ggml_tensor * pwconv2_w = nullptr; + struct ggml_tensor * pwconv2_b = nullptr; + struct ggml_tensor * gamma = nullptr; +}; + +// Audio tokenizer decoder model weights +struct audio_decoder_model { + audio_decoder_config config; + + // VQ codebooks + // vq_first: 1 codebook for first code + struct ggml_tensor * vq_first_input_proj = nullptr; // [1, 512, 256] + struct ggml_tensor * vq_first_output_proj = nullptr; // [1, 256, 512] + struct ggml_tensor * vq_first_codebook = nullptr; // [256, 2048] embedding_sum + struct ggml_tensor * vq_first_usage = nullptr; // [2048] cluster_usage + + // vq_rest: 15 codebooks for remaining codes + struct ggml_tensor * vq_rest_input_proj = nullptr; // [1, 512, 256] + struct ggml_tensor * vq_rest_output_proj = nullptr; // [1, 256, 512] + struct ggml_tensor * vq_rest_codebook[15] = {nullptr}; // [256, 2048] embedding_sum each + struct ggml_tensor * vq_rest_usage[15] = {nullptr}; // [2048] cluster_usage each + + // Upsample blocks (2 ConvNeXt-style blocks) + upsample_block upsample[2]; + + // Pre-transformer + struct ggml_tensor * pre_tfm_input_proj_w = nullptr; // [1024, 512] + struct ggml_tensor * pre_tfm_input_proj_b = nullptr; + pre_tfm_layer pre_tfm_layers[8]; + struct ggml_tensor * pre_tfm_norm_w = nullptr; // Final RMSNorm + struct ggml_tensor * pre_tfm_output_proj_w = nullptr; // [512, 1024] + struct ggml_tensor * pre_tfm_output_proj_b = nullptr; + + // Pre-conv: [3, 512, 1024] + struct ggml_tensor * pre_conv_w = nullptr; + struct ggml_tensor * pre_conv_b = nullptr; + + // Decoder blocks + // Block 0: Initial conv [7, 1024, 1536] + struct ggml_tensor * dec0_conv_w = nullptr; + struct ggml_tensor * dec0_conv_b = nullptr; + + // Blocks 1-4: Snake + ConvTranspose + 3 residual blocks + decoder_block dec_blocks[4]; + + // Block 5: Final snake activation + struct ggml_tensor * dec5_snake_alpha = nullptr; + struct ggml_tensor * dec5_snake_beta = nullptr; + + // Block 6: Output conv [7, 96, 1] + struct ggml_tensor * dec6_conv_w = nullptr; + struct ggml_tensor * dec6_conv_b = nullptr; + + // GGML context for tensor metadata + struct ggml_context * ctx = nullptr; + + // Backend buffer for weights + ggml_backend_buffer_t buffer = nullptr; + + // Tensor name to tensor mapping + std::map tensors; +}; + +// Compute state for decoder +struct audio_decoder_state { + ggml_backend_t backend = nullptr; + ggml_backend_t backend_cpu = nullptr; + ggml_backend_sched_t sched = nullptr; + std::vector compute_meta; +}; + +// Audio tokenizer decoder (vocoder) class +// Decodes discrete audio codes to waveform +class AudioTokenizerDecoder { +public: + AudioTokenizerDecoder(); + ~AudioTokenizerDecoder(); + + // Load model from GGUF file (tokenizer model) + bool load_model(const std::string & model_path); + + // Release all model/runtime resources + void unload_model(); + + // Decode audio codes to waveform + // codes: audio codes [n_frames, n_codebooks] as int32_t (row-major) + // n_frames: number of frames + // Returns: audio samples normalized to [-1, 1] at 24kHz + bool decode(const int32_t * codes, int32_t n_frames, + std::vector & samples); + + const audio_decoder_config & get_config() const { return model_.config; } + + const std::string & get_error() const { return error_msg_; } + +private: + // Build computation graph for decoding + struct ggml_cgraph * build_graph(int32_t n_frames); + + // Apply Snake activation: x + (1/alpha) * sin^2(alpha * x) + struct ggml_tensor * apply_snake(struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * alpha, + struct ggml_tensor * beta); + + // Apply RMSNorm + struct ggml_tensor * apply_rms_norm(struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * w, + float eps); + + // Apply pre-transformer layer + struct ggml_tensor * apply_pre_tfm_layer(struct ggml_context * ctx, + struct ggml_tensor * x, + const pre_tfm_layer & layer, + int32_t n_frames, + struct ggml_tensor * positions); + + // Apply upsample block (ConvNeXt-style) + struct ggml_tensor * apply_upsample_block(struct ggml_context * ctx, + struct ggml_tensor * x, + const upsample_block & block, + int block_idx); + + // Apply residual block + struct ggml_tensor * apply_residual_block(struct ggml_context * ctx, + struct ggml_tensor * x, + const residual_block & block); + + // Apply decoder block (Snake + ConvTranspose + Residuals) + struct ggml_tensor * apply_decoder_block(struct ggml_context * ctx, + struct ggml_tensor * x, + const decoder_block & block, + int upsample_rate, + int block_idx); + + void normalize_codebooks(); + + audio_decoder_model model_; + audio_decoder_state state_; + std::string error_msg_; + + // Temporary storage for codes input + std::vector codes_buf_; +}; + +// Free model resources +void free_audio_decoder_model(audio_decoder_model & model); + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/audio_tokenizer_encoder.cpp b/otherarch/qwen3tts/audio_tokenizer_encoder.cpp new file mode 100644 index 000000000..d2646737c --- /dev/null +++ b/otherarch/qwen3tts/audio_tokenizer_encoder.cpp @@ -0,0 +1,771 @@ +#include "audio_tokenizer_encoder.h" +#include "gguf_loader.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include + +#define QWEN3_TTS_MAX_NODES 16384 + +namespace qwen3_tts { + +// Mel filterbank computation using librosa slaney normalization +// This matches librosa.filters.mel with norm='slaney' +static void compute_mel_filterbank_slaney(float * filterbank, int n_mels, int n_fft, + int sample_rate, float f_min, float f_max) { + // Slaney-style mel scale (used by librosa default) + auto hz_to_mel_slaney = [](float hz) -> float { + // Linear below 1000 Hz, logarithmic above + const float f_sp = 200.0f / 3.0f; // 66.67 Hz + const float min_log_hz = 1000.0f; + const float min_log_mel = (min_log_hz - 0.0f) / f_sp; // 15 + const float logstep = logf(6.4f) / 27.0f; // log(6400/1000) / 27 + + if (hz < min_log_hz) { + return (hz - 0.0f) / f_sp; + } else { + return min_log_mel + logf(hz / min_log_hz) / logstep; + } + }; + + auto mel_to_hz_slaney = [](float mel) -> float { + const float f_sp = 200.0f / 3.0f; + const float min_log_hz = 1000.0f; + const float min_log_mel = (min_log_hz - 0.0f) / f_sp; + const float logstep = logf(6.4f) / 27.0f; + + if (mel < min_log_mel) { + return 0.0f + f_sp * mel; + } else { + return min_log_hz * expf(logstep * (mel - min_log_mel)); + } + }; + + float mel_min = hz_to_mel_slaney(f_min); + float mel_max = hz_to_mel_slaney(f_max); + + int n_fft_bins = n_fft / 2 + 1; + + // Compute mel center frequencies + std::vector mel_points(n_mels + 2); + for (int i = 0; i < n_mels + 2; ++i) { + mel_points[i] = mel_min + (mel_max - mel_min) * i / (n_mels + 1); + } + + // Convert to Hz and then to FFT bin indices + std::vector hz_points(n_mels + 2); + std::vector fft_freqs(n_fft_bins); + + for (int i = 0; i < n_mels + 2; ++i) { + hz_points[i] = mel_to_hz_slaney(mel_points[i]); + } + + for (int i = 0; i < n_fft_bins; ++i) { + fft_freqs[i] = (float)i * sample_rate / n_fft; + } + + memset(filterbank, 0, n_mels * n_fft_bins * sizeof(float)); + + // Create triangular filters with slaney normalization + for (int m = 0; m < n_mels; ++m) { + float f_left = hz_points[m]; + float f_center = hz_points[m + 1]; + float f_right = hz_points[m + 2]; + + // Slaney normalization: divide by bandwidth (area normalization) + float enorm = 2.0f / (f_right - f_left); + + for (int k = 0; k < n_fft_bins; ++k) { + float freq = fft_freqs[k]; + + if (freq >= f_left && freq <= f_center) { + if (f_center > f_left) { + filterbank[m * n_fft_bins + k] = enorm * (freq - f_left) / (f_center - f_left); + } + } else if (freq > f_center && freq <= f_right) { + if (f_right > f_center) { + filterbank[m * n_fft_bins + k] = enorm * (f_right - freq) / (f_right - f_center); + } + } + } + } +} + +static void compute_dft(const float * input, float * real, float * imag, int n) { + for (int k = 0; k < n; ++k) { + real[k] = 0.0f; + imag[k] = 0.0f; + for (int t = 0; t < n; ++t) { + float angle = -2.0f * M_PI * k * t / n; + real[k] += input[t] * cosf(angle); + imag[k] += input[t] * sinf(angle); + } + } +} + +// Periodic Hann window (matches torch.hann_window with periodic=True, which is default) +static void compute_hann_window(float * window, int n) { + for (int i = 0; i < n; ++i) { + window[i] = 0.5f * (1.0f - cosf(2.0f * M_PI * i / n)); + } +} + +// Compute centered window for STFT (PyTorch centers win_length window in n_fft frame) +static void compute_centered_window(float * window, int n_fft, int win_length) { + // Zero-initialize + memset(window, 0, n_fft * sizeof(float)); + + // Compute Hann window of win_length + int offset = (n_fft - win_length) / 2; + for (int i = 0; i < win_length; ++i) { + window[offset + i] = 0.5f * (1.0f - cosf(2.0f * M_PI * i / win_length)); + } +} + +AudioTokenizerEncoder::AudioTokenizerEncoder() = default; + +AudioTokenizerEncoder::~AudioTokenizerEncoder() { + free_speaker_encoder_model(model_); + + if (state_.sched) { + ggml_backend_sched_free(state_.sched); + state_.sched = nullptr; + } + if (state_.backend) { + release_preferred_backend(state_.backend); + state_.backend = nullptr; + } + if (state_.backend_cpu) { + ggml_backend_free(state_.backend_cpu); + state_.backend_cpu = nullptr; + } +} + +bool AudioTokenizerEncoder::load_model(const std::string & model_path) { + GGUFLoader loader; + if (!loader.open(model_path)) { + error_msg_ = loader.get_error(); + return false; + } + + model_.config.sample_rate = loader.get_u32("qwen3-tts.speaker_encoder.sample_rate", 24000); + model_.config.embedding_dim = loader.get_u32("qwen3-tts.speaker_encoder.embedding_length", 1024); + + int64_t n_tensors = loader.get_n_tensors(); + int spk_tensor_count = 0; + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = loader.get_tensor_name(i); + if (name && strncmp(name, "spk_enc.", 8) == 0) { + spk_tensor_count++; + } + } + + if (spk_tensor_count == 0) { + error_msg_ = "No speaker encoder tensors found in model"; + return false; + } + + size_t ctx_size = ggml_tensor_overhead() * spk_tensor_count; + struct ggml_init_params params = { + /*.mem_size =*/ ctx_size, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + model_.ctx = ggml_init(params); + if (!model_.ctx) { + error_msg_ = "Failed to initialize GGML context"; + return false; + } + + struct gguf_context * gguf_ctx = loader.get_ctx(); + struct ggml_context * meta_ctx = loader.get_meta_ctx(); + + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = loader.get_tensor_name(i); + if (!name || strncmp(name, "spk_enc.", 8) != 0) { + continue; + } + + struct ggml_tensor * meta_tensor = ggml_get_tensor(meta_ctx, name); + if (!meta_tensor) { + continue; + } + + struct ggml_tensor * tensor = ggml_dup_tensor(model_.ctx, meta_tensor); + ggml_set_name(tensor, name); + + model_.tensors[name] = tensor; + + std::string sname(name); + + if (sname == "spk_enc.conv0.weight") model_.conv0_w = tensor; + else if (sname == "spk_enc.conv0.bias") model_.conv0_b = tensor; + else if (sname == "spk_enc.mfa.weight") model_.mfa_w = tensor; + else if (sname == "spk_enc.mfa.bias") model_.mfa_b = tensor; + else if (sname == "spk_enc.asp.conv.weight") model_.asp_conv_w = tensor; + else if (sname == "spk_enc.asp.conv.bias") model_.asp_conv_b = tensor; + else if (sname == "spk_enc.asp.tdnn.weight") model_.asp_tdnn_w = tensor; + else if (sname == "spk_enc.asp.tdnn.bias") model_.asp_tdnn_b = tensor; + else if (sname == "spk_enc.fc.weight") model_.fc_w = tensor; + else if (sname == "spk_enc.fc.bias") model_.fc_b = tensor; + else { + int blk_idx, res_idx; + char suffix[64]; + + if (sscanf(name, "spk_enc.blk.%d.tdnn1.%s", &blk_idx, suffix) == 2) { + if (blk_idx >= 1 && blk_idx <= 3) { + if (strcmp(suffix, "weight") == 0) model_.blocks[blk_idx-1].tdnn1_w = tensor; + else if (strcmp(suffix, "bias") == 0) model_.blocks[blk_idx-1].tdnn1_b = tensor; + } + } + else if (sscanf(name, "spk_enc.blk.%d.tdnn2.%s", &blk_idx, suffix) == 2) { + if (blk_idx >= 1 && blk_idx <= 3) { + if (strcmp(suffix, "weight") == 0) model_.blocks[blk_idx-1].tdnn2_w = tensor; + else if (strcmp(suffix, "bias") == 0) model_.blocks[blk_idx-1].tdnn2_b = tensor; + } + } + else if (sscanf(name, "spk_enc.blk.%d.res2net.%d.%s", &blk_idx, &res_idx, suffix) == 3) { + if (blk_idx >= 1 && blk_idx <= 3 && res_idx >= 0 && res_idx < 7) { + if (strcmp(suffix, "weight") == 0) model_.blocks[blk_idx-1].res2net_w[res_idx] = tensor; + else if (strcmp(suffix, "bias") == 0) model_.blocks[blk_idx-1].res2net_b[res_idx] = tensor; + } + } + else if (sscanf(name, "spk_enc.blk.%d.se.conv1.%s", &blk_idx, suffix) == 2) { + if (blk_idx >= 1 && blk_idx <= 3) { + if (strcmp(suffix, "weight") == 0) model_.blocks[blk_idx-1].se_conv1_w = tensor; + else if (strcmp(suffix, "bias") == 0) model_.blocks[blk_idx-1].se_conv1_b = tensor; + } + } + else if (sscanf(name, "spk_enc.blk.%d.se.conv2.%s", &blk_idx, suffix) == 2) { + if (blk_idx >= 1 && blk_idx <= 3) { + if (strcmp(suffix, "weight") == 0) model_.blocks[blk_idx-1].se_conv2_w = tensor; + else if (strcmp(suffix, "bias") == 0) model_.blocks[blk_idx-1].se_conv2_b = tensor; + } + } + } + } + + if (!load_tensor_data_from_file(model_path, gguf_ctx, model_.ctx, + model_.tensors, model_.buffer, error_msg_)) { + return false; + } + + state_.backend = init_preferred_backend("AudioTokenizerEncoder", &error_msg_); + if (!state_.backend) { + return false; + } + ggml_backend_dev_t device = ggml_backend_get_device(state_.backend); + const char * device_name = device ? ggml_backend_dev_name(device) : "Unknown"; + fprintf(stderr, " AudioTokenizerEncoder backend: %s\n", device_name); + + if (device && ggml_backend_dev_type(device) != GGML_BACKEND_DEVICE_TYPE_CPU) { + state_.backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + if (!state_.backend_cpu) { + error_msg_ = "Failed to initialize CPU fallback backend for AudioTokenizerEncoder"; + return false; + } + } + + std::vector backends; + backends.push_back(state_.backend); + if (state_.backend_cpu) { + backends.push_back(state_.backend_cpu); + } + state_.sched = ggml_backend_sched_new(backends.data(), nullptr, (int)backends.size(), QWEN3_TTS_MAX_NODES, false, true); + if (!state_.sched) { + error_msg_ = "Failed to create backend scheduler"; + return false; + } + + state_.compute_meta.resize(ggml_tensor_overhead() * QWEN3_TTS_MAX_NODES + ggml_graph_overhead()); + + return true; +} + +bool AudioTokenizerEncoder::compute_mel_spectrogram(const float * samples, int32_t n_samples, + std::vector & mel, int32_t & n_frames) { + const auto & cfg = model_.config; + + // Match PyTorch STFT padding: (n_fft - hop_size) // 2 on each side with reflect + int padding = (cfg.n_fft - cfg.hop_length) / 2; + int padded_length = n_samples + 2 * padding; + + // Create padded signal with reflect padding + std::vector padded(padded_length); + for (int i = 0; i < padded_length; ++i) { + int src_idx; + if (i < padding) { + // Reflect left: padding-1, padding-2, ..., 0 -> samples[padding-i], samples[padding-1-i], ... + src_idx = padding - i; + } else if (i >= padding + n_samples) { + // Reflect right + src_idx = 2 * n_samples - (i - padding) - 2; + } else { + src_idx = i - padding; + } + // Clamp to valid range + src_idx = std::max(0, std::min(n_samples - 1, src_idx)); + padded[i] = samples[src_idx]; + } + + // With center=False, frames start at 0 and step by hop_length + n_frames = (padded_length - cfg.n_fft) / cfg.hop_length + 1; + if (n_frames <= 0) { + error_msg_ = "Audio too short for mel spectrogram"; + return false; + } + + int n_fft_bins = cfg.n_fft / 2 + 1; + + std::vector filterbank(cfg.n_mels * n_fft_bins); + compute_mel_filterbank_slaney(filterbank.data(), cfg.n_mels, cfg.n_fft, + cfg.sample_rate, cfg.f_min, cfg.f_max); + + // PyTorch STFT with win_length < n_fft centers the window in the n_fft frame + // This is critical for matching PyTorch's output + std::vector window(cfg.n_fft); + compute_centered_window(window.data(), cfg.n_fft, cfg.win_length); + + // Output: [batch, n_mels, n_frames] but we store as [n_mels, n_frames] row-major + // which means mel[m * n_frames + f] = value at mel bin m, frame f + mel.resize(cfg.n_mels * n_frames); + + std::vector frame(cfg.n_fft, 0.0f); + std::vector fft_real(cfg.n_fft); + std::vector fft_imag(cfg.n_fft); + std::vector magnitude(n_fft_bins); + + for (int32_t f = 0; f < n_frames; ++f) { + int start = f * cfg.hop_length; + + // Apply centered window to n_fft samples + for (int i = 0; i < cfg.n_fft; ++i) { + frame[i] = padded[start + i] * window[i]; + } + + compute_dft(frame.data(), fft_real.data(), fft_imag.data(), cfg.n_fft); + + // Compute magnitude (not power) - matches torch.stft with return_complex=True then abs() + // spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9) + for (int k = 0; k < n_fft_bins; ++k) { + magnitude[k] = sqrtf(fft_real[k] * fft_real[k] + fft_imag[k] * fft_imag[k] + 1e-9f); + } + + // Apply mel filterbank and log compression + // mel_spec = torch.matmul(mel_basis, spec) + // mel_spec = dynamic_range_compression_torch(mel_spec) # log(clamp(x, min=1e-5) * 1) + for (int m = 0; m < cfg.n_mels; ++m) { + float sum = 0.0f; + for (int k = 0; k < n_fft_bins; ++k) { + sum += filterbank[m * n_fft_bins + k] * magnitude[k]; + } + // dynamic_range_compression: log(clamp(x, min=1e-5)) + mel[m * n_frames + f] = logf(std::max(sum, 1e-5f)); + } + } + + return true; +} + +static struct ggml_tensor * apply_reflect_pad_1d(struct ggml_context * ctx, + struct ggml_tensor * x, + int pad) { + if (pad == 0) { + return x; + } + + int64_t T = x->ne[0]; + int64_t C = x->ne[1]; + int64_t B = x->ne[2]; + + struct ggml_tensor * left_slices[16]; + struct ggml_tensor * right_slices[16]; + + for (int i = 0; i < pad && i < 16; ++i) { + int left_src_idx = pad - i; + left_slices[i] = ggml_view_3d(ctx, x, 1, C, B, + x->nb[1], x->nb[2], + left_src_idx * x->nb[0]); + left_slices[i] = ggml_cont(ctx, left_slices[i]); + + int right_src_idx = T - 2 - i; + right_slices[i] = ggml_view_3d(ctx, x, 1, C, B, + x->nb[1], x->nb[2], + right_src_idx * x->nb[0]); + right_slices[i] = ggml_cont(ctx, right_slices[i]); + } + + struct ggml_tensor * left_pad = left_slices[0]; + for (int i = 1; i < pad && i < 16; ++i) { + left_pad = ggml_concat(ctx, left_pad, left_slices[i], 0); + } + + struct ggml_tensor * right_pad = right_slices[0]; + for (int i = 1; i < pad && i < 16; ++i) { + right_pad = ggml_concat(ctx, right_pad, right_slices[i], 0); + } + + struct ggml_tensor * padded = ggml_concat(ctx, left_pad, x, 0); + padded = ggml_concat(ctx, padded, right_pad, 0); + + return padded; +} + +static struct ggml_tensor * apply_conv1d(struct ggml_context * ctx, + struct ggml_tensor * w, + struct ggml_tensor * b, + struct ggml_tensor * x, + int stride, int pad, int dilation, + const char * debug_name = nullptr, + bool use_reflect_pad = true) { + struct ggml_tensor * input = x; + int actual_pad = pad; + + if (use_reflect_pad && pad > 0) { + input = apply_reflect_pad_1d(ctx, x, pad); + actual_pad = 0; + } + + struct ggml_tensor * y = ggml_conv_1d(ctx, w, input, stride, actual_pad, dilation); + if (debug_name) { + char name[64]; + snprintf(name, sizeof(name), "%s_conv", debug_name); + ggml_set_name(y, name); + } + if (b) { + int64_t oc = y->ne[1]; + y = ggml_add(ctx, y, ggml_reshape_3d(ctx, b, 1, oc, 1)); + } + return y; +} + +struct ggml_cgraph * AudioTokenizerEncoder::build_graph(int32_t n_frames) { + const auto & cfg = model_.config; + const int hidden_dim = cfg.hidden_dim; // 512 + const int scale = cfg.res2net_scale; // 8 + const int branch_dim = hidden_dim / scale; // 64 + + struct ggml_init_params params = { + /*.mem_size =*/ state_.compute_meta.size(), + /*.mem_buffer =*/ state_.compute_meta.data(), + /*.no_alloc =*/ true, + }; + + struct ggml_context * ctx0 = ggml_init(params); + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false); + + // Input: mel spectrogram [n_mels, n_frames] - stored as [n_mels, n_frames] row-major + // GGML uses column-major, so this is [n_frames, n_mels] in GGML notation + struct ggml_tensor * mel = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_frames, cfg.n_mels); + ggml_set_name(mel, "mel"); + ggml_set_input(mel); + + // PyTorch: hidden_states = hidden_states.transpose(1, 2) # [B, T, C] -> [B, C, T] + // Our mel is [n_frames, n_mels] in GGML = [n_mels, n_frames] row-major + // For conv1d, we need [T, C, B] in GGML = [B, C, T] row-major + // So reshape to [n_frames, n_mels, 1] + struct ggml_tensor * cur = ggml_reshape_3d(ctx0, mel, n_frames, cfg.n_mels, 1); + ggml_set_name(cur, "mel_3d"); + + struct ggml_tensor * mel_padded = apply_reflect_pad_1d(ctx0, cur, 2); + ggml_set_name(mel_padded, "mel_padded"); + + cur = ggml_conv_1d(ctx0, model_.conv0_w, mel_padded, 1, 0, 1); + ggml_set_name(cur, "conv0_conv"); + + if (model_.conv0_b) { + int64_t oc = cur->ne[1]; + cur = ggml_add(ctx0, cur, ggml_reshape_3d(ctx0, model_.conv0_b, 1, oc, 1)); + } + ggml_set_name(cur, "conv0_pre_relu"); + cur = ggml_relu(ctx0, cur); + ggml_set_name(cur, "conv0_out"); + + int64_t seq_len = cur->ne[0]; + + // Store block outputs for MFA (including block 0) + struct ggml_tensor * block_outputs[4]; + block_outputs[0] = cur; // Initial TDNN output + + // Blocks 1-3: SE-Res2Net blocks + // Dilations: block1=2, block2=3, block3=4 + int dilations[3] = {2, 3, 4}; + + for (int blk = 0; blk < 3; ++blk) { + const auto & block = model_.blocks[blk]; + int dilation = dilations[blk]; + + struct ggml_tensor * residual = cur; + + cur = apply_conv1d(ctx0, block.tdnn1_w, block.tdnn1_b, cur, 1, 0, 1); + cur = ggml_relu(ctx0, cur); + if (blk == 0) { + ggml_set_name(cur, "blk1_tdnn1"); + } + + // Res2Net: Split into 8 branches of 64 channels each + // cur shape: [seq_len, 512, 1] + // Branch 0: identity (no conv) + // Branch i (1-7): conv(hidden_part + previous_output) for i >= 2, conv(hidden_part) for i == 1 + + // Split channels: view as [seq_len, 64, 8] then split + struct ggml_tensor * branches[8]; + + // Extract each branch using view operations + // cur is [seq_len, 512, 1], we want to split dim 1 into 8 parts of 64 + for (int b = 0; b < scale; ++b) { + // View into the b-th chunk of 64 channels + // cur shape: [seq_len, 512, 1], we want [seq_len, 64, 1] starting at channel b*64 + // nb1 = stride for dim 1 = cur->nb[1] (bytes to move from one channel to next) + // nb2 = stride for dim 2 = cur->nb[2] (bytes to move from one batch to next) + // offset = b * 64 * cur->nb[1] (skip b*64 channels) + branches[b] = ggml_view_3d(ctx0, cur, + seq_len, branch_dim, 1, + cur->nb[1], cur->nb[2], + b * branch_dim * cur->nb[1]); + branches[b] = ggml_cont(ctx0, branches[b]); + } + + // Process branches according to Res2Net logic + struct ggml_tensor * outputs[8]; + outputs[0] = branches[0]; // Branch 0: identity + + for (int b = 1; b < scale; ++b) { + struct ggml_tensor * input; + if (b == 1) { + input = branches[b]; + } else { + // Add previous output to current branch + input = ggml_add(ctx0, branches[b], outputs[b - 1]); + } + + // Apply conv with dilation (kernel=3) + // Padding for kernel=3, dilation=d: pad = d * (3-1) / 2 = d + if (block.res2net_w[b - 1]) { + outputs[b] = apply_conv1d(ctx0, block.res2net_w[b - 1], block.res2net_b[b - 1], + input, 1, dilation, dilation); + outputs[b] = ggml_relu(ctx0, outputs[b]); + } else { + outputs[b] = input; // Fallback if weight missing + } + } + + cur = outputs[0]; + for (int b = 1; b < scale; ++b) { + cur = ggml_concat(ctx0, cur, outputs[b], 1); + } + if (blk == 0) { + ggml_set_name(cur, "blk1_res2net"); + for (int b = 0; b < scale; ++b) { + char name[32]; + snprintf(name, sizeof(name), "blk1_branch%d", b); + ggml_set_name(outputs[b], name); + } + } + + cur = apply_conv1d(ctx0, block.tdnn2_w, block.tdnn2_b, cur, 1, 0, 1); + cur = ggml_relu(ctx0, cur); + if (blk == 0) { + ggml_set_name(cur, "blk1_tdnn2"); + } + + // SE (Squeeze-Excitation) + // Global average pooling over time: mean(dim=2, keepdim=True) + struct ggml_tensor * se = ggml_pool_1d(ctx0, cur, GGML_OP_POOL_AVG, seq_len, seq_len, 0); + se = ggml_reshape_3d(ctx0, se, 1, hidden_dim, 1); + + // SE conv1: 512 -> 128 with ReLU + se = apply_conv1d(ctx0, block.se_conv1_w, block.se_conv1_b, se, 1, 0, 1); + se = ggml_relu(ctx0, se); + + // SE conv2: 128 -> 512 with Sigmoid + se = apply_conv1d(ctx0, block.se_conv2_w, block.se_conv2_b, se, 1, 0, 1); + se = ggml_sigmoid(ctx0, se); + + cur = ggml_mul(ctx0, cur, se); + if (blk == 0) { + ggml_set_name(cur, "blk1_se"); + } + + cur = ggml_add(ctx0, cur, residual); + + char block_name[32]; + snprintf(block_name, sizeof(block_name), "block_%d", blk + 1); + ggml_set_name(cur, block_name); + + block_outputs[blk + 1] = cur; + } + + // MFA: Concatenate block outputs [1:] (blocks 1, 2, 3 = indices 1, 2, 3) + // hidden_states = torch.cat(hidden_states_list[1:], dim=1) + // Each block output is [seq_len, 512, 1] + // Concatenated: [seq_len, 1536, 1] + struct ggml_tensor * mfa_input = ggml_concat(ctx0, block_outputs[1], block_outputs[2], 1); + mfa_input = ggml_concat(ctx0, mfa_input, block_outputs[3], 1); + ggml_set_name(mfa_input, "mfa_input"); + + // MFA conv: 1536 -> 1536 with ReLU + cur = apply_conv1d(ctx0, model_.mfa_w, model_.mfa_b, mfa_input, 1, 0, 1); + cur = ggml_relu(ctx0, cur); + ggml_set_name(cur, "mfa_out"); + + // ASP (Attentive Statistics Pooling) + // cur shape: [seq_len, 1536, 1] + + // Step 1: Compute global mean and std over time + // mean = hidden_states.mean(dim=2, keepdim=True) # [1, 1536, 1] + struct ggml_tensor * global_mean = ggml_pool_1d(ctx0, cur, GGML_OP_POOL_AVG, seq_len, seq_len, 0); + global_mean = ggml_reshape_3d(ctx0, global_mean, 1, 1536, 1); + + // std = sqrt(E[x^2] - E[x]^2) + struct ggml_tensor * sq = ggml_sqr(ctx0, cur); + struct ggml_tensor * mean_sq = ggml_pool_1d(ctx0, sq, GGML_OP_POOL_AVG, seq_len, seq_len, 0); + mean_sq = ggml_reshape_3d(ctx0, mean_sq, 1, 1536, 1); + struct ggml_tensor * var = ggml_sub(ctx0, mean_sq, ggml_sqr(ctx0, global_mean)); + var = ggml_clamp(ctx0, var, 1e-12f, 1e10f); + struct ggml_tensor * global_std = ggml_sqrt(ctx0, var); + + // Step 2: Expand mean and std to full sequence length and concatenate with hidden_states + // mean = mean.repeat(1, 1, seq_length) # [1, 1536, seq_len] + // std = std.repeat(1, 1, seq_length) # [1, 1536, seq_len] + // attention = torch.cat([hidden_states, mean, std], dim=1) # [1, 4608, seq_len] + struct ggml_tensor * mean_expanded = ggml_repeat(ctx0, global_mean, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, seq_len, 1536, 1)); + struct ggml_tensor * std_expanded = ggml_repeat(ctx0, global_std, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, seq_len, 1536, 1)); + + struct ggml_tensor * attention = ggml_concat(ctx0, cur, mean_expanded, 1); + attention = ggml_concat(ctx0, attention, std_expanded, 1); + // attention shape: [seq_len, 4608, 1] + + // Step 3: TDNN (4608 -> 128) with ReLU, then Tanh + // self.tdnn = TimeDelayNetBlock(channels * 3, attention_channels, 1, 1) # has ReLU + // attention = self.conv(self.tanh(self.tdnn(attention))) + attention = apply_conv1d(ctx0, model_.asp_tdnn_w, model_.asp_tdnn_b, attention, 1, 0, 1); + attention = ggml_relu(ctx0, attention); // TDNN has ReLU + ggml_set_name(attention, "asp_tdnn"); + attention = ggml_tanh(ctx0, attention); // Then tanh is applied + + // Step 4: Conv (128 -> 1536) for attention weights + // self.conv = nn.Conv1d(attention_channels, channels, kernel_size=1) + attention = apply_conv1d(ctx0, model_.asp_conv_w, model_.asp_conv_b, attention, 1, 0, 1); + ggml_set_name(attention, "asp_conv"); + // attention shape: [seq_len, 1536, 1] + + // Step 5: Softmax over time dimension + attention = ggml_soft_max(ctx0, attention); + ggml_set_name(attention, "asp_softmax"); + + // Step 6: Compute weighted mean and std + // mean, std = self._compute_statistics(hidden_states, attention) + // mean = (attention * hidden_states).sum(dim=2) + struct ggml_tensor * weighted = ggml_mul(ctx0, attention, cur); + struct ggml_tensor * weighted_mean = ggml_pool_1d(ctx0, weighted, GGML_OP_POOL_AVG, seq_len, seq_len, 0); + weighted_mean = ggml_scale(ctx0, weighted_mean, (float)seq_len); // Convert avg to sum + weighted_mean = ggml_reshape_3d(ctx0, weighted_mean, 1, 1536, 1); + + // std = sqrt((attention * (hidden_states - mean)^2).sum(dim=2).clamp(eps)) + struct ggml_tensor * mean_for_std = ggml_repeat(ctx0, weighted_mean, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, seq_len, 1536, 1)); + struct ggml_tensor * diff = ggml_sub(ctx0, cur, mean_for_std); + struct ggml_tensor * diff_sq = ggml_sqr(ctx0, diff); + struct ggml_tensor * weighted_var = ggml_mul(ctx0, attention, diff_sq); + struct ggml_tensor * var_sum = ggml_pool_1d(ctx0, weighted_var, GGML_OP_POOL_AVG, seq_len, seq_len, 0); + var_sum = ggml_scale(ctx0, var_sum, (float)seq_len); // Convert avg to sum + var_sum = ggml_reshape_3d(ctx0, var_sum, 1, 1536, 1); + var_sum = ggml_clamp(ctx0, var_sum, 1e-12f, 1e10f); + struct ggml_tensor * weighted_std = ggml_sqrt(ctx0, var_sum); + + // Step 7: Concatenate mean and std: [1, 3072, 1] + struct ggml_tensor * pooled = ggml_concat(ctx0, weighted_mean, weighted_std, 1); + ggml_set_name(pooled, "asp_pooled"); + + // FC: 3072 -> 1024 + cur = apply_conv1d(ctx0, model_.fc_w, model_.fc_b, pooled, 1, 0, 1); + ggml_set_name(cur, "fc_out"); + + // Squeeze to 1D + cur = ggml_reshape_1d(ctx0, cur, cfg.embedding_dim); + + ggml_set_name(cur, "embedding"); + ggml_set_output(cur); + + ggml_build_forward_expand(gf, cur); + + ggml_free(ctx0); + + return gf; +} + +bool AudioTokenizerEncoder::encode(const float * samples, int32_t n_samples, + std::vector & embedding) { + if (!model_.ctx) { + error_msg_ = "Model not loaded"; + return false; + } + + std::vector mel; + int32_t n_frames; + if (!compute_mel_spectrogram(samples, n_samples, mel, n_frames)) { + return false; + } + + struct ggml_cgraph * gf = build_graph(n_frames); + + if (!ggml_backend_sched_alloc_graph(state_.sched, gf)) { + error_msg_ = "Failed to allocate graph"; + return false; + } + + struct ggml_tensor * mel_tensor = ggml_graph_get_tensor(gf, "mel"); + if (!mel_tensor) { + error_msg_ = "Failed to find mel tensor"; + ggml_backend_sched_reset(state_.sched); + return false; + } + + // mel is stored as [n_mels, n_frames] row-major: mel[m * n_frames + f] = mel bin m at frame f + // GGML tensor is [n_frames, n_mels] column-major: element (f, m) at memory[f + m * n_frames] + // For GGML conv1d, we want input(t, c) = mel bin c at time t + // So GGML memory[t + c * n_frames] should equal mel[c * n_frames + t] + // Since the memory layout matches (both are contiguous in frame order for each mel bin), + // we can copy directly! + ggml_backend_tensor_set(mel_tensor, mel.data(), 0, mel.size() * sizeof(float)); + + if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) { + error_msg_ = "Failed to compute graph"; + ggml_backend_sched_reset(state_.sched); + return false; + } + + struct ggml_tensor * emb_tensor = ggml_graph_get_tensor(gf, "embedding"); + if (!emb_tensor) { + error_msg_ = "Failed to find embedding tensor"; + ggml_backend_sched_reset(state_.sched); + return false; + } + + embedding.resize(model_.config.embedding_dim); + ggml_backend_tensor_get(emb_tensor, embedding.data(), 0, embedding.size() * sizeof(float)); + + ggml_backend_sched_reset(state_.sched); + + return true; +} + +void free_speaker_encoder_model(speaker_encoder_model & model) { + if (model.buffer) { + ggml_backend_buffer_free(model.buffer); + model.buffer = nullptr; + } + if (model.ctx) { + ggml_free(model.ctx); + model.ctx = nullptr; + } + model.tensors.clear(); +} + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/audio_tokenizer_encoder.h b/otherarch/qwen3tts/audio_tokenizer_encoder.h new file mode 100644 index 000000000..6dc856bd8 --- /dev/null +++ b/otherarch/qwen3tts/audio_tokenizer_encoder.h @@ -0,0 +1,154 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" +#include "gguf.h" + +#include +#include +#include +#include + +namespace qwen3_tts { + +// Speaker encoder configuration (ECAPA-TDNN) +// Mel parameters MUST match extract_speaker_embedding() in modeling_qwen3_tts.py +struct speaker_encoder_config { + int32_t sample_rate = 24000; + int32_t n_mels = 128; + int32_t n_fft = 1024; + int32_t hop_length = 256; + int32_t win_length = 1024; + int32_t embedding_dim = 1024; + int32_t hidden_dim = 512; + int32_t n_res2net_blocks = 3; + int32_t res2net_scale = 8; + float f_min = 0.0f; + float f_max = 12000.0f; +}; + +// Res2Net block weights +struct res2net_block { + // TDNN1: 1x1 conv (512 -> 512) + struct ggml_tensor * tdnn1_w = nullptr; + struct ggml_tensor * tdnn1_b = nullptr; + + // Res2Net branches: 7 conv layers (kernel=3, 64 -> 64) + struct ggml_tensor * res2net_w[7] = {nullptr}; + struct ggml_tensor * res2net_b[7] = {nullptr}; + + // TDNN2: 1x1 conv (512 -> 512) + struct ggml_tensor * tdnn2_w = nullptr; + struct ggml_tensor * tdnn2_b = nullptr; + + // SE (Squeeze-Excitation) + struct ggml_tensor * se_conv1_w = nullptr; + struct ggml_tensor * se_conv1_b = nullptr; + struct ggml_tensor * se_conv2_w = nullptr; + struct ggml_tensor * se_conv2_b = nullptr; +}; + +// Speaker encoder model weights +struct speaker_encoder_model { + speaker_encoder_config config; + + // Initial conv: (5, 128, 512) - kernel 5, in 128 (mel), out 512 + struct ggml_tensor * conv0_w = nullptr; + struct ggml_tensor * conv0_b = nullptr; + + // Res2Net blocks (3 blocks) + res2net_block blocks[3]; + + // MFA (Multi-Frame Aggregation): 1x1 conv (1536 -> 1536) + struct ggml_tensor * mfa_w = nullptr; + struct ggml_tensor * mfa_b = nullptr; + + // ASP (Attentive Statistics Pooling) + struct ggml_tensor * asp_conv_w = nullptr; + struct ggml_tensor * asp_conv_b = nullptr; + struct ggml_tensor * asp_tdnn_w = nullptr; + struct ggml_tensor * asp_tdnn_b = nullptr; + + // Final FC: (3072 -> 1024) + struct ggml_tensor * fc_w = nullptr; + struct ggml_tensor * fc_b = nullptr; + + // GGML context for tensor metadata + struct ggml_context * ctx = nullptr; + + // Backend buffer for weights + ggml_backend_buffer_t buffer = nullptr; + + // Tensor name to tensor mapping + std::map tensors; +}; + +// Compute state for speaker encoder +struct speaker_encoder_state { + ggml_backend_t backend = nullptr; + ggml_backend_t backend_cpu = nullptr; + ggml_backend_sched_t sched = nullptr; + std::vector compute_meta; +}; + +// Speaker encoder class (ECAPA-TDNN) +// Extracts speaker embedding from audio waveform +class AudioTokenizerEncoder { +public: + AudioTokenizerEncoder(); + ~AudioTokenizerEncoder(); + + // Load model from GGUF file (main TTS model, not tokenizer) + bool load_model(const std::string & model_path); + + // Encode audio samples to speaker embedding + // samples: audio samples normalized to [-1, 1], 24kHz + // n_samples: number of samples + // embedding: output speaker embedding [1024] + bool encode(const float * samples, int32_t n_samples, + std::vector & embedding); + + // Legacy interface for compatibility (not used for speaker encoding) + bool encode(const float * samples, int32_t n_samples, + std::vector & codes, int32_t & n_frames) { + (void)samples; (void)n_samples; (void)codes; (void)n_frames; + error_msg_ = "Use encode(samples, n_samples, embedding) for speaker encoding"; + return false; + } + + // Legacy interface (not used) + bool get_embeddings(const int32_t * codes, int32_t n_frames, + std::vector & embeddings) { + (void)codes; (void)n_frames; (void)embeddings; + error_msg_ = "Use encode() for speaker embedding extraction"; + return false; + } + + const speaker_encoder_config & get_config() const { return model_.config; } + + const std::string & get_error() const { return error_msg_; } + +private: + // Compute mel spectrogram from waveform + bool compute_mel_spectrogram(const float * samples, int32_t n_samples, + std::vector & mel, int32_t & n_frames); + + // Build computation graph + struct ggml_cgraph * build_graph(int32_t n_frames); + + speaker_encoder_model model_; + speaker_encoder_state state_; + std::string error_msg_; +}; + +// Free model resources +void free_speaker_encoder_model(speaker_encoder_model & model); + +// Backward compatibility alias +using audio_encoder_config = speaker_encoder_config; +using audio_encoder_model = speaker_encoder_model; +inline void free_audio_encoder_model(audio_encoder_model & model) { + free_speaker_encoder_model(model); +} + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/coreml_code_predictor.h b/otherarch/qwen3tts/coreml_code_predictor.h new file mode 100644 index 000000000..cff351a83 --- /dev/null +++ b/otherarch/qwen3tts/coreml_code_predictor.h @@ -0,0 +1,35 @@ +#pragma once + +#include +#include +#include + +namespace qwen3_tts { + +class CoreMLCodePredictor { +public: + CoreMLCodePredictor(); + ~CoreMLCodePredictor(); + + CoreMLCodePredictor(const CoreMLCodePredictor &) = delete; + CoreMLCodePredictor & operator=(const CoreMLCodePredictor &) = delete; + + bool load(const std::string & model_dir, int32_t n_steps); + void unload(); + + bool is_loaded() const; + const std::string & get_error() const; + + bool predict_step(int32_t step_idx, + const float * seq_embd, + int32_t seq_len, + int32_t hidden_size, + std::vector & logits_out); + +private: + struct Impl; + Impl * impl_ = nullptr; +}; + +} // namespace qwen3_tts + diff --git a/otherarch/qwen3tts/coreml_code_predictor_stub.cpp b/otherarch/qwen3tts/coreml_code_predictor_stub.cpp new file mode 100644 index 000000000..7b0b3d977 --- /dev/null +++ b/otherarch/qwen3tts/coreml_code_predictor_stub.cpp @@ -0,0 +1,31 @@ +#include "coreml_code_predictor.h" + +namespace qwen3_tts { + +CoreMLCodePredictor::CoreMLCodePredictor() {} +CoreMLCodePredictor::~CoreMLCodePredictor() {} + +bool CoreMLCodePredictor::load(const std::string &, int32_t) { + return false; +} + +void CoreMLCodePredictor::unload() {} + +bool CoreMLCodePredictor::is_loaded() const { + return false; +} + +const std::string & CoreMLCodePredictor::get_error() const { + static const std::string err = "CoreML predictor only supported on Apple platforms"; + return err; +} + +bool CoreMLCodePredictor::predict_step(int32_t, + const float *, + int32_t, + int32_t, + std::vector &) { + return false; +} + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/gguf_loader.cpp b/otherarch/qwen3tts/gguf_loader.cpp new file mode 100644 index 000000000..649af520d --- /dev/null +++ b/otherarch/qwen3tts/gguf_loader.cpp @@ -0,0 +1,238 @@ +#include "gguf_loader.h" + +#include +#include +#include + +namespace qwen3_tts { + +namespace { +struct shared_backend_state { + ggml_backend_t backend = nullptr; + int32_t ref_count = 0; +}; + +shared_backend_state & get_shared_backend_state() { + static shared_backend_state state; + return state; +} +} + +GGUFLoader::GGUFLoader() = default; + +GGUFLoader::~GGUFLoader() { + close(); +} + +ggml_backend_t init_preferred_backend(const char * component_name, std::string * error_msg) { + if (error_msg) error_msg->clear(); + + auto & shared = get_shared_backend_state(); + if (shared.backend) { + shared.ref_count++; + return shared.backend; + } + + ggml_backend_t backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU, nullptr); + if (!backend) { + backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_GPU, nullptr); + } + if (!backend) { + backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_ACCEL, nullptr); + } + if (!backend) { + backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + } + + if (!backend && error_msg) { + const char * name = component_name ? component_name : "component"; + *error_msg = "Failed to initialize backend (IGPU/GPU/ACCEL/CPU) for " + std::string(name); + } + + if (backend) { + shared.backend = backend; + shared.ref_count = 1; + } + + return backend; +} + +void release_preferred_backend(ggml_backend_t backend) { + if (!backend) { + return; + } + + auto & shared = get_shared_backend_state(); + if (shared.backend == backend) { + shared.ref_count--; + if (shared.ref_count <= 0) { + ggml_backend_free(shared.backend); + shared.backend = nullptr; + shared.ref_count = 0; + } + return; + } + + ggml_backend_free(backend); +} + +bool GGUFLoader::open(const std::string & path) { + close(); // Close any previously opened file + + file_path_ = path; + + struct gguf_init_params params = { + /*.no_alloc =*/ true, + /*.ctx =*/ &meta_ctx_, + }; + + ctx_ = gguf_init_from_file(path.c_str(), params); + if (!ctx_) { + error_msg_ = "Failed to open GGUF file: " + path; + return false; + } + + return true; +} + +void GGUFLoader::close() { + if (ctx_) { + gguf_free(ctx_); + ctx_ = nullptr; + } + if (meta_ctx_) { + ggml_free(meta_ctx_); + meta_ctx_ = nullptr; + } + file_path_.clear(); +} + +int64_t GGUFLoader::get_n_tensors() const { + if (!ctx_) return 0; + return gguf_get_n_tensors(ctx_); +} + +const char * GGUFLoader::get_tensor_name(int64_t idx) const { + if (!ctx_) return nullptr; + return gguf_get_tensor_name(ctx_, idx); +} + +enum ggml_type GGUFLoader::get_tensor_type(int64_t idx) const { + if (!ctx_) return GGML_TYPE_F32; + return gguf_get_tensor_type(ctx_, idx); +} + +size_t GGUFLoader::get_tensor_offset(int64_t idx) const { + if (!ctx_) return 0; + return gguf_get_tensor_offset(ctx_, idx); +} + +size_t GGUFLoader::get_tensor_size(int64_t idx) const { + if (!ctx_) return 0; + return gguf_get_tensor_size(ctx_, idx); +} + +int32_t GGUFLoader::get_u32(const char * key, int32_t default_val) const { + if (!ctx_) return default_val; + int64_t idx = gguf_find_key(ctx_, key); + if (idx < 0) return default_val; + return (int32_t)gguf_get_val_u32(ctx_, idx); +} + +float GGUFLoader::get_f32(const char * key, float default_val) const { + if (!ctx_) return default_val; + int64_t idx = gguf_find_key(ctx_, key); + if (idx < 0) return default_val; + return gguf_get_val_f32(ctx_, idx); +} + +size_t GGUFLoader::get_data_offset() const { + if (!ctx_) return 0; + return gguf_get_data_offset(ctx_); +} + +bool load_tensor_data_from_file( + const std::string & path, + struct gguf_context * ctx, + struct ggml_context * model_ctx, + const std::map & tensors, + ggml_backend_buffer_t & buffer, + std::string & error_msg, + enum ggml_backend_dev_type preferred_backend_type +) { + ggml_backend_t backend = ggml_backend_init_by_type(preferred_backend_type, nullptr); + if (!backend && preferred_backend_type != GGML_BACKEND_DEVICE_TYPE_CPU) { + backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + } + if (!backend) { + error_msg = "Failed to initialize backend for GGUF tensor loader"; + return false; + } + + // Allocate buffer for all tensors + buffer = ggml_backend_alloc_ctx_tensors(model_ctx, backend); + if (!buffer) { + error_msg = "Failed to allocate tensor buffer"; + ggml_backend_free(backend); + return false; + } + + // Open file for reading tensor data + FILE * f = fopen(path.c_str(), "rb"); + if (!f) { + error_msg = "Failed to open file for reading: " + path; + ggml_backend_free(backend); + return false; + } + + const size_t data_offset = gguf_get_data_offset(ctx); + const int64_t n_tensors = gguf_get_n_tensors(ctx); + std::vector read_buf; + + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = gguf_get_tensor_name(ctx, i); + size_t offset = gguf_get_tensor_offset(ctx, i); + + auto it = tensors.find(name); + if (it == tensors.end()) { + continue; // Skip tensors not in our map + } + + struct ggml_tensor * tensor = it->second; + size_t nbytes = ggml_nbytes(tensor); + + read_buf.resize(nbytes); + + if (fseek(f, data_offset + offset, SEEK_SET) != 0) { + error_msg = "Failed to seek to tensor data: " + std::string(name); + fclose(f); + ggml_backend_free(backend); + return false; + } + + if (fread(read_buf.data(), 1, nbytes, f) != nbytes) { + error_msg = "Failed to read tensor data: " + std::string(name); + fclose(f); + ggml_backend_free(backend); + return false; + } + + ggml_backend_tensor_set(tensor, read_buf.data(), 0, nbytes); + } + + fclose(f); + ggml_backend_free(backend); + + return true; +} + +void free_ggml_resources(struct ggml_context * ctx, ggml_backend_buffer_t buffer) { + if (buffer) { + ggml_backend_buffer_free(buffer); + } + if (ctx) { + ggml_free(ctx); + } +} + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/gguf_loader.h b/otherarch/qwen3tts/gguf_loader.h new file mode 100644 index 000000000..d14d84f9c --- /dev/null +++ b/otherarch/qwen3tts/gguf_loader.h @@ -0,0 +1,83 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" +#include "gguf.h" + +#include +#include +#include +#include + +namespace qwen3_tts { + +// Generic GGUF model loader class +// This is a simplified loader that can be extended for specific model types +class GGUFLoader { +public: + GGUFLoader(); + ~GGUFLoader(); + + // Open GGUF file and parse metadata + bool open(const std::string & path); + + // Close file and free resources + void close(); + + // Get error message if operation failed + const std::string & get_error() const { return error_msg_; } + + // Get number of tensors in file + int64_t get_n_tensors() const; + + // Get tensor name by index + const char * get_tensor_name(int64_t idx) const; + + // Get tensor type by index + enum ggml_type get_tensor_type(int64_t idx) const; + + // Get tensor offset by index + size_t get_tensor_offset(int64_t idx) const; + + // Get tensor size by index + size_t get_tensor_size(int64_t idx) const; + + // Get metadata value (returns -1 if not found) + int32_t get_u32(const char * key, int32_t default_val = 0) const; + float get_f32(const char * key, float default_val = 0.0f) const; + + // Get data offset (start of tensor data in file) + size_t get_data_offset() const; + + // Get GGUF context (for advanced usage) + struct gguf_context * get_ctx() const { return ctx_; } + + // Get metadata context + struct ggml_context * get_meta_ctx() const { return meta_ctx_; } + +protected: + struct gguf_context * ctx_ = nullptr; + struct ggml_context * meta_ctx_ = nullptr; + std::string error_msg_; + std::string file_path_; +}; + +// Helper function to allocate and load tensor data from GGUF file +bool load_tensor_data_from_file( + const std::string & path, + struct gguf_context * ctx, + struct ggml_context * model_ctx, + const std::map & tensors, + ggml_backend_buffer_t & buffer, + std::string & error_msg, + enum ggml_backend_dev_type preferred_backend_type = GGML_BACKEND_DEVICE_TYPE_CPU +); + +// Helper to initialize backend with GPU preference and CPU fallback +ggml_backend_t init_preferred_backend(const char * component_name, std::string * error_msg); +void release_preferred_backend(ggml_backend_t backend); + +// Helper function to free model resources +void free_ggml_resources(struct ggml_context * ctx, ggml_backend_buffer_t buffer); + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/main.cpp b/otherarch/qwen3tts/main.cpp new file mode 100644 index 000000000..77f81ce0e --- /dev/null +++ b/otherarch/qwen3tts/main.cpp @@ -0,0 +1,178 @@ +#include "qwen3_tts.h" + +#include +#include +#include + +void print_usage(const char * program) { + fprintf(stderr, "Usage: %s [options] -m -t \n", program); + fprintf(stderr, "\n"); + fprintf(stderr, "Options:\n"); + fprintf(stderr, " -m, --model Model directory (required)\n"); + fprintf(stderr, " -t, --text Text to synthesize (required)\n"); + fprintf(stderr, " -o, --output Output WAV file (default: output.wav)\n"); + fprintf(stderr, " -r, --reference Reference audio for voice cloning\n"); + fprintf(stderr, " --temperature Sampling temperature (default: 0.9, 0=greedy)\n"); + fprintf(stderr, " --top-k Top-k sampling (default: 50, 0=disabled)\n"); + fprintf(stderr, " --top-p Top-p sampling (default: 1.0)\n"); + fprintf(stderr, " --max-tokens Maximum audio tokens (default: 4096)\n"); + fprintf(stderr, " --repetition-penalty Repetition penalty (default: 1.05)\n"); + fprintf(stderr, " -j, --threads Number of threads (default: 4)\n"); + fprintf(stderr, " -h, --help Show this help\n"); + fprintf(stderr, "\n"); + fprintf(stderr, "Example:\n"); + fprintf(stderr, " %s -m ./models -t \"Hello, world!\" -o hello.wav\n", program); + fprintf(stderr, " %s -m ./models -t \"Hello!\" -r reference.wav -o cloned.wav\n", program); +} + +int main(int argc, char ** argv) { + std::string model_dir; + std::string text; + std::string output_file = "output.wav"; + std::string reference_audio; + + qwen3_tts::tts_params params; + + // Parse arguments + for (int i = 1; i < argc; i++) { + std::string arg = argv[i]; + + if (arg == "-h" || arg == "--help") { + print_usage(argv[0]); + return 0; + } else if (arg == "-m" || arg == "--model") { + if (++i >= argc) { + fprintf(stderr, "Error: missing model directory\n"); + return 1; + } + model_dir = argv[i]; + } else if (arg == "-t" || arg == "--text") { + if (++i >= argc) { + fprintf(stderr, "Error: missing text\n"); + return 1; + } + text = argv[i]; + } else if (arg == "-o" || arg == "--output") { + if (++i >= argc) { + fprintf(stderr, "Error: missing output file\n"); + return 1; + } + output_file = argv[i]; + } else if (arg == "-r" || arg == "--reference") { + if (++i >= argc) { + fprintf(stderr, "Error: missing reference audio\n"); + return 1; + } + reference_audio = argv[i]; + } else if (arg == "--temperature") { + if (++i >= argc) { + fprintf(stderr, "Error: missing temperature value\n"); + return 1; + } + params.temperature = std::stof(argv[i]); + } else if (arg == "--top-k") { + if (++i >= argc) { + fprintf(stderr, "Error: missing top-k value\n"); + return 1; + } + params.top_k = std::stoi(argv[i]); + } else if (arg == "--top-p") { + if (++i >= argc) { + fprintf(stderr, "Error: missing top-p value\n"); + return 1; + } + params.top_p = std::stof(argv[i]); + } else if (arg == "--max-tokens") { + if (++i >= argc) { + fprintf(stderr, "Error: missing max-tokens value\n"); + return 1; + } + params.max_audio_tokens = std::stoi(argv[i]); + } else if (arg == "--repetition-penalty") { + if (++i >= argc) { + fprintf(stderr, "Error: missing repetition-penalty value\n"); + return 1; + } + params.repetition_penalty = std::stof(argv[i]); + } else if (arg == "-j" || arg == "--threads") { + if (++i >= argc) { + fprintf(stderr, "Error: missing threads value\n"); + return 1; + } + params.n_threads = std::stoi(argv[i]); + } else { + fprintf(stderr, "Error: unknown argument: %s\n", arg.c_str()); + print_usage(argv[0]); + return 1; + } + } + + // Validate required arguments + if (model_dir.empty()) { + fprintf(stderr, "Error: model directory is required\n"); + print_usage(argv[0]); + return 1; + } + + if (text.empty()) { + fprintf(stderr, "Error: text is required\n"); + print_usage(argv[0]); + return 1; + } + + // Initialize TTS + qwen3_tts::Qwen3TTS tts; + + fprintf(stderr, "Loading models from: %s\n", model_dir.c_str()); + if (!tts.load_models(model_dir)) { + fprintf(stderr, "Error: %s\n", tts.get_error().c_str()); + return 1; + } + + // Set progress callback + tts.set_progress_callback([](int tokens, int max_tokens) { + fprintf(stderr, "\rGenerating: %d/%d tokens", tokens, max_tokens); + }); + + // Generate speech + qwen3_tts::tts_result result; + + if (reference_audio.empty()) { + fprintf(stderr, "Synthesizing: \"%s\"\n", text.c_str()); + result = tts.synthesize(text, params); + } else { + fprintf(stderr, "Synthesizing with voice cloning: \"%s\"\n", text.c_str()); + fprintf(stderr, "Reference audio: %s\n", reference_audio.c_str()); + result = tts.synthesize_with_voice(text, reference_audio, params); + } + + if (!result.success) { + fprintf(stderr, "\nError: %s\n", result.error_msg.c_str()); + return 1; + } + + fprintf(stderr, "\n"); + + // Save output + if (!qwen3_tts::save_audio_file(output_file, result.audio, result.sample_rate)) { + fprintf(stderr, "Error: failed to save output file: %s\n", output_file.c_str()); + return 1; + } + + fprintf(stderr, "Output saved to: %s\n", output_file.c_str()); + fprintf(stderr, "Audio duration: %.2f seconds\n", + (float)result.audio.size() / result.sample_rate); + + // Print timing + if (params.print_timing) { + fprintf(stderr, "\nTiming:\n"); + fprintf(stderr, " Load: %6lld ms\n", (long long)result.t_load_ms); + fprintf(stderr, " Tokenize: %6lld ms\n", (long long)result.t_tokenize_ms); + fprintf(stderr, " Encode: %6lld ms\n", (long long)result.t_encode_ms); + fprintf(stderr, " Generate: %6lld ms\n", (long long)result.t_generate_ms); + fprintf(stderr, " Decode: %6lld ms\n", (long long)result.t_decode_ms); + fprintf(stderr, " Total: %6lld ms\n", (long long)result.t_total_ms); + } + + return 0; +} diff --git a/otherarch/qwen3tts/qwen3_tts.cpp b/otherarch/qwen3tts/qwen3_tts.cpp new file mode 100644 index 000000000..07b540226 --- /dev/null +++ b/otherarch/qwen3tts/qwen3_tts.cpp @@ -0,0 +1,585 @@ +#include "qwen3_tts.h" +#include "gguf_loader.h" + +#include +#include +#include +#include +#include +#include +#include + + + +namespace qwen3_tts { + +static int64_t get_time_ms() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +struct process_memory_snapshot { + uint64_t rss_bytes = 0; + uint64_t phys_footprint_bytes = 0; +}; + +static bool get_process_memory_snapshot(process_memory_snapshot & out) { +return false; +} + +static std::string format_bytes(uint64_t bytes) { + static const char * units[] = { "B", "KB", "MB", "GB", "TB" }; + double val = (double) bytes; + int unit = 0; + while (val >= 1024.0 && unit < 4) { + val /= 1024.0; + ++unit; + } + char buf[64]; + snprintf(buf, sizeof(buf), "%.2f %s", val, units[unit]); + return std::string(buf); +} + +static void log_memory_usage(const char * label) { + +} + +static void resample_linear(const float * input, int input_len, int input_rate, + std::vector & output, int output_rate) { + double ratio = (double)input_rate / output_rate; + int output_len = (int)((double)input_len / ratio); + output.resize(output_len); + + for (int i = 0; i < output_len; ++i) { + double src_idx = i * ratio; + int idx0 = (int)src_idx; + int idx1 = idx0 + 1; + double frac = src_idx - idx0; + + if (idx1 >= input_len) { + output[i] = input[input_len - 1]; + } else { + output[i] = (float)((1.0 - frac) * input[idx0] + frac * input[idx1]); + } + } +} + +Qwen3TTS::Qwen3TTS() = default; + +Qwen3TTS::~Qwen3TTS() = default; + +bool Qwen3TTS::load_models(const std::string & model_dir) { + int64_t t_start = get_time_ms(); + log_memory_usage("load/start"); + + transformer_.unload_model(); + audio_decoder_.unload_model(); + transformer_loaded_ = false; + decoder_loaded_ = false; + + // Construct model paths + std::string tts_model_path = model_dir + "/qwen3-tts-0.6b-f16.gguf"; + std::string tokenizer_model_path = model_dir + "/qwen3-tts-tokenizer-f16.gguf"; + tts_model_path_ = tts_model_path; + decoder_model_path_ = tokenizer_model_path; + encoder_loaded_ = false; + transformer_loaded_ = false; + decoder_loaded_ = false; + + const char * low_mem_env = std::getenv("QWEN3_TTS_LOW_MEM"); + low_mem_mode_ = low_mem_env && low_mem_env[0] != '\0' && low_mem_env[0] != '0'; + if (low_mem_mode_) { + fprintf(stderr, " Low-memory mode enabled (lazy decoder + component unloads)\n"); + } + + // Load TTS model (contains text tokenizer + transformer for generation) + fprintf(stderr, "Loading TTS model from %s...\n", tts_model_path.c_str()); + + // Load text tokenizer from TTS model + int64_t t_tokenizer_start = get_time_ms(); + { + GGUFLoader loader; + if (!loader.open(tts_model_path)) { + error_msg_ = "Failed to open TTS model: " + loader.get_error(); + return false; + } + + if (!tokenizer_.load_from_gguf(loader.get_ctx())) { + error_msg_ = "Failed to load text tokenizer: " + tokenizer_.get_error(); + return false; + } + fprintf(stderr, " Text tokenizer loaded: vocab_size=%d (%lld ms)\n", + tokenizer_.get_config().vocab_size, + (long long)(get_time_ms() - t_tokenizer_start)); + } + log_memory_usage("load/after-tokenizer"); + + // Speaker encoder is loaded lazily on first voice cloning request. + fprintf(stderr, " Speaker encoder: deferred (lazy load)\n"); + + // Load TTS transformer from TTS model + int64_t t_transformer_start = get_time_ms(); + if (!transformer_.load_model(tts_model_path)) { + error_msg_ = "Failed to load TTS transformer: " + transformer_.get_error(); + return false; + } + transformer_loaded_ = true; + fprintf(stderr, " TTS transformer loaded: hidden_size=%d, n_layers=%d (%lld ms)\n", + transformer_.get_config().hidden_size, transformer_.get_config().n_layers, + (long long)(get_time_ms() - t_transformer_start)); + log_memory_usage("load/after-transformer"); + + if (!low_mem_mode_) { + // Load vocoder (audio decoder) from tokenizer model + fprintf(stderr, "Loading vocoder from %s...\n", tokenizer_model_path.c_str()); + int64_t t_decoder_start = get_time_ms(); + if (!audio_decoder_.load_model(tokenizer_model_path)) { + error_msg_ = "Failed to load vocoder: " + audio_decoder_.get_error(); + return false; + } + decoder_loaded_ = true; + fprintf(stderr, " Vocoder loaded: sample_rate=%d, n_codebooks=%d (%lld ms)\n", + audio_decoder_.get_config().sample_rate, audio_decoder_.get_config().n_codebooks, + (long long)(get_time_ms() - t_decoder_start)); + log_memory_usage("load/after-vocoder"); + } else { + fprintf(stderr, " Vocoder: deferred (lazy load)\n"); + } + + models_loaded_ = true; + + int64_t t_end = get_time_ms(); + fprintf(stderr, "All models loaded in %lld ms\n", (long long)(t_end - t_start)); + log_memory_usage("load/end"); + + return true; +} + +tts_result Qwen3TTS::synthesize(const std::string & text, + const tts_params & params) { + tts_result result; + + if (!models_loaded_) { + result.error_msg = "Models not loaded"; + return result; + } + + // For basic synthesis without voice cloning, we use a zero speaker embedding + // This will use the model's default voice characteristics + std::vector zero_embedding(transformer_.get_config().hidden_size, 0.0f); + + return synthesize_internal(text, zero_embedding.data(), params, result); +} + +tts_result Qwen3TTS::synthesize_with_voice(const std::string & text, + const std::string & reference_audio, + const tts_params & params) { + tts_result result; + + std::vector ref_samples; + int ref_sample_rate; + if (!load_audio_file(reference_audio, ref_samples, ref_sample_rate)) { + result.error_msg = "Failed to load reference audio: " + reference_audio; + return result; + } + + const int target_rate = 24000; + if (ref_sample_rate != target_rate) { + fprintf(stderr, "Resampling audio from %d Hz to %d Hz...\n", ref_sample_rate, target_rate); + std::vector resampled; + resample_linear(ref_samples.data(), (int)ref_samples.size(), ref_sample_rate, resampled, target_rate); + ref_samples = std::move(resampled); + } + + return synthesize_with_voice(text, ref_samples.data(), (int32_t)ref_samples.size(), params); +} + +tts_result Qwen3TTS::synthesize_with_voice(const std::string & text, + const float * ref_samples, int32_t n_ref_samples, + const tts_params & params) { + tts_result result; + + if (!models_loaded_) { + result.error_msg = "Models not loaded"; + return result; + } + + if (!encoder_loaded_) { + if (tts_model_path_.empty()) { + result.error_msg = "Internal error: missing TTS model path for lazy encoder load"; + return result; + } + int64_t t_encoder_load_start = get_time_ms(); + if (!audio_encoder_.load_model(tts_model_path_)) { + result.error_msg = "Failed to load speaker encoder: " + audio_encoder_.get_error(); + return result; + } + encoder_loaded_ = true; + if (params.print_timing) { + fprintf(stderr, " Speaker encoder lazy-loaded in %lld ms\n", + (long long)(get_time_ms() - t_encoder_load_start)); + log_memory_usage("voice/after-encoder-load"); + } + } + + int64_t t_encode_start = get_time_ms(); + std::vector speaker_embedding; + + if (!audio_encoder_.encode(ref_samples, n_ref_samples, speaker_embedding)) { + result.error_msg = "Failed to extract speaker embedding: " + audio_encoder_.get_error(); + return result; + } + result.t_encode_ms = get_time_ms() - t_encode_start; + + if (params.print_progress) { + fprintf(stderr, "Speaker embedding extracted: %zu floats\n", speaker_embedding.size()); + } + + return synthesize_internal(text, speaker_embedding.data(), params, result); +} + +tts_result Qwen3TTS::synthesize_internal(const std::string & text, + const float * speaker_embedding, + const tts_params & params, + tts_result & result) { + int64_t t_total_start = get_time_ms(); + auto sample_memory = [&](const char * stage) { + + }; + sample_memory("synth/start"); + + // Step 2: Tokenize input text + int64_t t_tokenize_start = get_time_ms(); + std::vector text_tokens = tokenizer_.encode_for_tts(text); + result.t_tokenize_ms = get_time_ms() - t_tokenize_start; + sample_memory("synth/after-tokenize"); + + if (text_tokens.empty()) { + result.error_msg = "Failed to tokenize text"; + return result; + } + + if (params.print_progress) { + fprintf(stderr, "Text tokenized: %zu tokens\n", text_tokens.size()); + fprintf(stderr, " Tokens: "); + for (size_t i = 0; i < std::min(text_tokens.size(), (size_t)10); ++i) { + fprintf(stderr, "%d ", text_tokens[i]); + } + if (text_tokens.size() > 10) fprintf(stderr, "..."); + fprintf(stderr, "\n"); + } + + // Step 3: Generate speech codes using TTS transformer + int64_t t_generate_start = get_time_ms(); + if (!transformer_loaded_) { + int64_t t_reload_start = get_time_ms(); + if (!transformer_.load_model(tts_model_path_)) { + result.error_msg = "Failed to reload TTS transformer: " + transformer_.get_error(); + return result; + } + transformer_loaded_ = true; + if (params.print_timing) { + fprintf(stderr, " Transformer reloaded in %lld ms\n", + (long long)(get_time_ms() - t_reload_start)); + sample_memory("synth/after-transformer-reload"); + } + } + transformer_.clear_kv_cache(); + + std::vector speech_codes; + if (!transformer_.generate(text_tokens.data(), (int32_t)text_tokens.size(), + speaker_embedding, params.max_audio_tokens, speech_codes, + 2050, params.repetition_penalty, + params.temperature, params.top_k)) { + result.error_msg = "Failed to generate speech codes: " + transformer_.get_error(); + return result; + } + result.t_generate_ms = get_time_ms() - t_generate_start; + sample_memory("synth/after-generate"); + + int n_codebooks = transformer_.get_config().n_codebooks; + int n_frames = (int)speech_codes.size() / n_codebooks; + + if (params.print_progress) { + fprintf(stderr, "Speech codes generated: %d frames x %d codebooks\n", n_frames, n_codebooks); + } + + if (n_frames == 0) { + result.error_msg = "No speech codes generated"; + return result; + } + + if (low_mem_mode_) { + transformer_.unload_model(); + transformer_loaded_ = false; + sample_memory("synth/after-transformer-unload"); + } + + // Step 4: Decode speech codes to waveform using vocoder + int64_t t_decode_start = get_time_ms(); + if (!decoder_loaded_) { + int64_t t_decoder_load_start = get_time_ms(); + if (decoder_model_path_.empty()) { + result.error_msg = "Internal error: missing vocoder model path"; + return result; + } + if (!audio_decoder_.load_model(decoder_model_path_)) { + result.error_msg = "Failed to load vocoder: " + audio_decoder_.get_error(); + return result; + } + decoder_loaded_ = true; + if (params.print_timing) { + fprintf(stderr, " Vocoder lazy-loaded in %lld ms\n", + (long long)(get_time_ms() - t_decoder_load_start)); + sample_memory("synth/after-vocoder-load"); + } + } + + if (!audio_decoder_.decode(speech_codes.data(), n_frames, result.audio)) { + result.error_msg = "Failed to decode speech codes: " + audio_decoder_.get_error(); + return result; + } + result.t_decode_ms = get_time_ms() - t_decode_start; + sample_memory("synth/after-decode"); + + if (low_mem_mode_) { + audio_decoder_.unload_model(); + decoder_loaded_ = false; + sample_memory("synth/after-vocoder-unload"); + } + + result.sample_rate = audio_decoder_.get_config().sample_rate; + result.success = true; + result.t_total_ms = get_time_ms() - t_total_start; + sample_memory("synth/end"); + + if (params.print_timing) { + const double audio_sec = result.sample_rate > 0 + ? (double) result.audio.size() / (double) result.sample_rate : 0.0; + const double wall_sec = (double) result.t_total_ms / 1000.0; + const double realtime_factor = audio_sec > 0.0 ? wall_sec / audio_sec : 0.0; + const double x_realtime = wall_sec > 0.0 ? audio_sec / wall_sec : 0.0; + fprintf(stderr, "\nTiming:\n"); + fprintf(stderr, " Tokenization: %lld ms\n", (long long)result.t_tokenize_ms); + fprintf(stderr, " Speaker encode: %lld ms\n", (long long)result.t_encode_ms); + fprintf(stderr, " Code generation: %lld ms\n", (long long)result.t_generate_ms); + fprintf(stderr, " Vocoder decode: %lld ms\n", (long long)result.t_decode_ms); + fprintf(stderr, " Total: %lld ms\n", (long long)result.t_total_ms); + fprintf(stderr, " Audio duration: %.2f s\n", audio_sec); + fprintf(stderr, " Throughput: %.2fx realtime (RTF=%.3f)\n", x_realtime, realtime_factor); + fprintf(stderr, "\nMemory:\n"); + fprintf(stderr, " RSS start/end: %s -> %s\n", + format_bytes(result.mem_rss_start_bytes).c_str(), + format_bytes(result.mem_rss_end_bytes).c_str()); + fprintf(stderr, " RSS peak: %s\n", + format_bytes(result.mem_rss_peak_bytes).c_str()); + fprintf(stderr, " Phys start/end: %s -> %s\n", + format_bytes(result.mem_phys_start_bytes).c_str(), + format_bytes(result.mem_phys_end_bytes).c_str()); + fprintf(stderr, " Phys peak: %s\n", + format_bytes(result.mem_phys_peak_bytes).c_str()); + } + + return result; +} + +void Qwen3TTS::set_progress_callback(tts_progress_callback_t callback) { + progress_callback_ = callback; +} + +// WAV file loading (16-bit PCM or 32-bit float) +bool load_audio_file(const std::string & path, std::vector & samples, + int & sample_rate) { + FILE * f = fopen(path.c_str(), "rb"); + if (!f) { + fprintf(stderr, "ERROR: Cannot open WAV file: %s\n", path.c_str()); + return false; + } + + // Read RIFF header + char riff[4]; + if (fread(riff, 1, 4, f) != 4 || strncmp(riff, "RIFF", 4) != 0) { + fprintf(stderr, "ERROR: Not a RIFF file\n"); + fclose(f); + return false; + } + + uint32_t file_size; + if (fread(&file_size, 4, 1, f) != 1) { + fclose(f); + return false; + } + + char wave[4]; + if (fread(wave, 1, 4, f) != 4 || strncmp(wave, "WAVE", 4) != 0) { + fprintf(stderr, "ERROR: Not a WAVE file\n"); + fclose(f); + return false; + } + + // Find fmt and data chunks + uint16_t audio_format = 0; + uint16_t num_channels = 0; + uint32_t sr = 0; + uint16_t bits_per_sample = 0; + + while (!feof(f)) { + char chunk_id[4]; + uint32_t chunk_size; + + if (fread(chunk_id, 1, 4, f) != 4) break; + if (fread(&chunk_size, 4, 1, f) != 1) break; + + if (strncmp(chunk_id, "fmt ", 4) == 0) { + if (fread(&audio_format, 2, 1, f) != 1) break; + if (fread(&num_channels, 2, 1, f) != 1) break; + if (fread(&sr, 4, 1, f) != 1) break; + fseek(f, 6, SEEK_CUR); // Skip byte rate and block align + if (fread(&bits_per_sample, 2, 1, f) != 1) break; + + // Skip any extra format bytes + if (chunk_size > 16) { + fseek(f, chunk_size - 16, SEEK_CUR); + } + } + else if (strncmp(chunk_id, "data", 4) == 0) { + sample_rate = sr; + + if (audio_format == 1) { // PCM + if (bits_per_sample == 16) { + int n_samples = chunk_size / (2 * num_channels); + samples.resize(n_samples); + + std::vector raw(n_samples * num_channels); + if (fread(raw.data(), 2, n_samples * num_channels, f) != (size_t)(n_samples * num_channels)) { + fclose(f); + return false; + } + + // Convert to mono float + for (int i = 0; i < n_samples; ++i) { + float sum = 0.0f; + for (int c = 0; c < num_channels; ++c) { + sum += raw[i * num_channels + c] / 32768.0f; + } + samples[i] = sum / num_channels; + } + } + else if (bits_per_sample == 32) { + int n_samples = chunk_size / (4 * num_channels); + samples.resize(n_samples); + + std::vector raw(n_samples * num_channels); + if (fread(raw.data(), 4, n_samples * num_channels, f) != (size_t)(n_samples * num_channels)) { + fclose(f); + return false; + } + + // Convert to mono float + for (int i = 0; i < n_samples; ++i) { + float sum = 0.0f; + for (int c = 0; c < num_channels; ++c) { + sum += raw[i * num_channels + c] / 2147483648.0f; + } + samples[i] = sum / num_channels; + } + } + else { + fprintf(stderr, "ERROR: Unsupported bits per sample: %d\n", bits_per_sample); + fclose(f); + return false; + } + } + else if (audio_format == 3) { // IEEE float + int n_samples = chunk_size / (4 * num_channels); + samples.resize(n_samples); + + std::vector raw(n_samples * num_channels); + if (fread(raw.data(), 4, n_samples * num_channels, f) != (size_t)(n_samples * num_channels)) { + fclose(f); + return false; + } + + // Convert to mono + for (int i = 0; i < n_samples; ++i) { + float sum = 0.0f; + for (int c = 0; c < num_channels; ++c) { + sum += raw[i * num_channels + c]; + } + samples[i] = sum / num_channels; + } + } + else { + fprintf(stderr, "ERROR: Unsupported audio format: %d\n", audio_format); + fclose(f); + return false; + } + + fclose(f); + return true; + } + else { + // Skip unknown chunk + fseek(f, chunk_size, SEEK_CUR); + } + } + + fprintf(stderr, "ERROR: No data chunk found\n"); + fclose(f); + return false; +} + +// WAV file saving (16-bit PCM at specified sample rate) +bool save_audio_file(const std::string & path, const std::vector & samples, + int sample_rate) { + FILE * f = fopen(path.c_str(), "wb"); + if (!f) { + fprintf(stderr, "ERROR: Cannot create WAV file: %s\n", path.c_str()); + return false; + } + + // WAV header parameters + uint16_t num_channels = 1; + uint16_t bits_per_sample = 16; + uint32_t byte_rate = sample_rate * num_channels * bits_per_sample / 8; + uint16_t block_align = num_channels * bits_per_sample / 8; + uint32_t data_size = samples.size() * block_align; + uint32_t file_size = 36 + data_size; + + // Write RIFF header + fwrite("RIFF", 1, 4, f); + fwrite(&file_size, 4, 1, f); + fwrite("WAVE", 1, 4, f); + + // Write fmt chunk + fwrite("fmt ", 1, 4, f); + uint32_t fmt_size = 16; + fwrite(&fmt_size, 4, 1, f); + uint16_t audio_format = 1; // PCM + fwrite(&audio_format, 2, 1, f); + fwrite(&num_channels, 2, 1, f); + uint32_t sr = sample_rate; + fwrite(&sr, 4, 1, f); + fwrite(&byte_rate, 4, 1, f); + fwrite(&block_align, 2, 1, f); + fwrite(&bits_per_sample, 2, 1, f); + + // Write data chunk + fwrite("data", 1, 4, f); + fwrite(&data_size, 4, 1, f); + + // Convert float samples to 16-bit PCM and write + for (size_t i = 0; i < samples.size(); ++i) { + // Clamp to [-1, 1] and convert to int16 + float sample = samples[i]; + if (sample > 1.0f) sample = 1.0f; + if (sample < -1.0f) sample = -1.0f; + int16_t pcm_sample = (int16_t)(sample * 32767.0f); + fwrite(&pcm_sample, 2, 1, f); + } + + fclose(f); + return true; +} + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/qwen3_tts.h b/otherarch/qwen3tts/qwen3_tts.h new file mode 100644 index 000000000..9008a4301 --- /dev/null +++ b/otherarch/qwen3tts/qwen3_tts.h @@ -0,0 +1,150 @@ +#pragma once + +#include "text_tokenizer.h" +#include "tts_transformer.h" +#include "audio_tokenizer_encoder.h" +#include "audio_tokenizer_decoder.h" + +#include +#include +#include +#include + +namespace qwen3_tts { + +// TTS generation parameters +struct tts_params { + // Maximum number of audio tokens to generate + int32_t max_audio_tokens = 4096; + + // Temperature for sampling (0 = greedy) + float temperature = 0.9f; + + // Top-p sampling + float top_p = 1.0f; + + // Top-k sampling (0 = disabled) + int32_t top_k = 50; + + // Number of threads + int32_t n_threads = 4; + + // Print progress during generation + bool print_progress = false; + + // Print timing information + bool print_timing = true; + + // Repetition penalty for CB0 token generation (HuggingFace style) + float repetition_penalty = 1.05f; + +}; + +// TTS generation result +struct tts_result { + // Generated audio samples (24kHz, mono) + std::vector audio; + + // Sample rate + int32_t sample_rate = 24000; + + // Success flag + bool success = false; + + // Error message if failed + std::string error_msg; + + // Timing info (in milliseconds) + int64_t t_load_ms = 0; + int64_t t_tokenize_ms = 0; + int64_t t_encode_ms = 0; + int64_t t_generate_ms = 0; + int64_t t_decode_ms = 0; + int64_t t_total_ms = 0; + + // Process memory snapshots (bytes) + uint64_t mem_rss_start_bytes = 0; + uint64_t mem_rss_end_bytes = 0; + uint64_t mem_rss_peak_bytes = 0; + uint64_t mem_phys_start_bytes = 0; + uint64_t mem_phys_end_bytes = 0; + uint64_t mem_phys_peak_bytes = 0; + +}; + +// Progress callback type +using tts_progress_callback_t = std::function; + +// Main TTS class that orchestrates the full pipeline +class Qwen3TTS { +public: + Qwen3TTS(); + ~Qwen3TTS(); + + // Load all models from directory + // model_dir should contain: transformer.gguf, tokenizer.gguf, vocoder.gguf + bool load_models(const std::string & model_dir); + + // Generate speech from text + // text: input text to synthesize + // params: generation parameters + tts_result synthesize(const std::string & text, + const tts_params & params = tts_params()); + + // Generate speech with voice cloning + // text: input text to synthesize + // reference_audio: path to reference audio file (WAV, 24kHz) + // params: generation parameters + tts_result synthesize_with_voice(const std::string & text, + const std::string & reference_audio, + const tts_params & params = tts_params()); + + // Generate speech with voice cloning from samples + // text: input text to synthesize + // ref_samples: reference audio samples (24kHz, mono, normalized to [-1, 1]) + // n_ref_samples: number of reference samples + // params: generation parameters + tts_result synthesize_with_voice(const std::string & text, + const float * ref_samples, int32_t n_ref_samples, + const tts_params & params = tts_params()); + + // Set progress callback + void set_progress_callback(tts_progress_callback_t callback); + + // Get error message + const std::string & get_error() const { return error_msg_; } + + // Check if models are loaded + bool is_loaded() const { return models_loaded_; } + +private: + tts_result synthesize_internal(const std::string & text, + const float * speaker_embedding, + const tts_params & params, + tts_result & result); + + TextTokenizer tokenizer_; + TTSTransformer transformer_; + AudioTokenizerEncoder audio_encoder_; + AudioTokenizerDecoder audio_decoder_; + + bool models_loaded_ = false; + bool encoder_loaded_ = false; + bool transformer_loaded_ = false; + bool decoder_loaded_ = false; + bool low_mem_mode_ = false; + std::string error_msg_; + std::string tts_model_path_; + std::string decoder_model_path_; + tts_progress_callback_t progress_callback_; +}; + +// Utility: Load audio file (WAV format) +bool load_audio_file(const std::string & path, std::vector & samples, + int & sample_rate); + +// Utility: Save audio file (WAV format) +bool save_audio_file(const std::string & path, const std::vector & samples, + int sample_rate); + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/text_tokenizer.cpp b/otherarch/qwen3tts/text_tokenizer.cpp new file mode 100644 index 000000000..d3a4aa9ee --- /dev/null +++ b/otherarch/qwen3tts/text_tokenizer.cpp @@ -0,0 +1,351 @@ +#include "text_tokenizer.h" + +#include +#include +#include +#include + +namespace qwen3_tts { + +// GPT-2 byte-to-unicode mapping +// Maps bytes 0-255 to unicode characters to avoid control characters +static const char * BYTE_TO_UNICODE[256] = { + "Ā", "ā", "Ă", "ă", "Ą", "ą", "Ć", "ć", "Ĉ", "ĉ", "Ċ", "ċ", "Č", "č", "Ď", "ď", + "Đ", "đ", "Ē", "ē", "Ĕ", "ĕ", "Ė", "ė", "Ę", "ę", "Ě", "ě", "Ĝ", "ĝ", "Ğ", "ğ", + "Ġ", "!", "\"", "#", "$", "%", "&", "'", "(", ")", "*", "+", ",", "-", ".", "/", + "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", ":", ";", "<", "=", ">", "?", + "@", "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", + "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "[", "\\", "]", "^", "_", + "`", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", + "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "{", "|", "}", "~", "ġ", + "Ģ", "ģ", "Ĥ", "ĥ", "Ħ", "ħ", "Ĩ", "ĩ", "Ī", "ī", "Ĭ", "ĭ", "Į", "į", "İ", "ı", + "IJ", "ij", "Ĵ", "ĵ", "Ķ", "ķ", "ĸ", "Ĺ", "ĺ", "Ļ", "ļ", "Ľ", "ľ", "Ŀ", "ŀ", "Ł", + "ł", "¡", "¢", "£", "¤", "¥", "¦", "§", "¨", "©", "ª", "«", "¬", "®", "¯", "°", + "±", "²", "³", "´", "µ", "¶", "·", "¸", "¹", "º", "»", "¼", "½", "¾", "¿", "À", + "Á", "Â", "Ã", "Ä", "Å", "Æ", "Ç", "È", "É", "Ê", "Ë", "Ì", "Í", "Î", "Ï", "Ð", + "Ñ", "Ò", "Ó", "Ô", "Õ", "Ö", "×", "Ø", "Ù", "Ú", "Û", "Ü", "Ý", "Þ", "ß", "à", + "á", "â", "ã", "ä", "å", "æ", "ç", "è", "é", "ê", "ë", "ì", "í", "î", "ï", "ð", + "ñ", "ò", "ó", "ô", "õ", "ö", "÷", "ø", "ù", "ú", "û", "ü", "ý", "þ", "ÿ", "Ń" +}; + +// Build reverse mapping at runtime +static std::unordered_map build_unicode_to_byte() { + std::unordered_map result; + for (int i = 0; i < 256; i++) { + result[BYTE_TO_UNICODE[i]] = (uint8_t)i; + } + return result; +} + +static const std::unordered_map UNICODE_TO_BYTE = build_unicode_to_byte(); + +TextTokenizer::TextTokenizer() = default; + +TextTokenizer::~TextTokenizer() = default; + +size_t TextTokenizer::utf8_len(char c) { + if ((c & 0x80) == 0) return 1; + if ((c & 0xE0) == 0xC0) return 2; + if ((c & 0xF0) == 0xE0) return 3; + if ((c & 0xF8) == 0xF0) return 4; + return 1; // Invalid UTF-8, treat as single byte +} + +std::string TextTokenizer::bytes_to_unicode(const std::string & text) { + std::string result; + for (unsigned char c : text) { + result += BYTE_TO_UNICODE[c]; + } + return result; +} + +std::string TextTokenizer::unicode_to_bytes(const std::string & text) { + std::string result; + size_t i = 0; + while (i < text.size()) { + size_t len = utf8_len(text[i]); + std::string ch = text.substr(i, len); + auto it = UNICODE_TO_BYTE.find(ch); + if (it != UNICODE_TO_BYTE.end()) { + result += (char)it->second; + } else { + // Not in mapping, keep as-is (shouldn't happen for valid tokens) + result += ch; + } + i += len; + } + return result; +} + +bool TextTokenizer::load_from_gguf(struct gguf_context * ctx) { + if (!ctx) { + error_msg_ = "GGUF context is null"; + return false; + } + + // Get vocabulary + int64_t tokens_key = gguf_find_key(ctx, "tokenizer.ggml.tokens"); + if (tokens_key < 0) { + error_msg_ = "tokenizer.ggml.tokens not found in GGUF"; + return false; + } + + size_t n_vocab = gguf_get_arr_n(ctx, tokens_key); + if (n_vocab == 0) { + error_msg_ = "Empty vocabulary"; + return false; + } + + config_.vocab_size = (int32_t)n_vocab; + id_to_token_.resize(n_vocab); + + for (size_t i = 0; i < n_vocab; i++) { + const char * token = gguf_get_arr_str(ctx, tokens_key, i); + if (token) { + id_to_token_[i] = token; + vocab_[token] = (int32_t)i; + } + } + + // Get merges + int64_t merges_key = gguf_find_key(ctx, "tokenizer.ggml.merges"); + if (merges_key >= 0) { + size_t n_merges = gguf_get_arr_n(ctx, merges_key); + for (size_t i = 0; i < n_merges; i++) { + const char * merge = gguf_get_arr_str(ctx, merges_key, i); + if (merge) { + std::string merge_str(merge); + // Parse "token1 token2" format + size_t space_pos = merge_str.find(' '); + if (space_pos != std::string::npos) { + std::string first = merge_str.substr(0, space_pos); + std::string second = merge_str.substr(space_pos + 1); + bpe_ranks_[{first, second}] = (int32_t)i; + } + } + } + } + + // Get special token IDs (optional, use defaults if not found) + int64_t bos_key = gguf_find_key(ctx, "tokenizer.ggml.bos_token_id"); + if (bos_key >= 0) { + config_.bos_token_id = (int32_t)gguf_get_val_u32(ctx, bos_key); + } + + int64_t eos_key = gguf_find_key(ctx, "tokenizer.ggml.eos_token_id"); + if (eos_key >= 0) { + config_.eos_token_id = (int32_t)gguf_get_val_u32(ctx, eos_key); + } + + int64_t pad_key = gguf_find_key(ctx, "tokenizer.ggml.padding_token_id"); + if (pad_key >= 0) { + config_.pad_token_id = (int32_t)gguf_get_val_u32(ctx, pad_key); + } + + // Find special tokens by content + auto find_token = [this](const std::string & text) -> int32_t { + auto it = vocab_.find(text); + return (it != vocab_.end()) ? it->second : -1; + }; + + assistant_token_id_ = find_token("assistant"); + if (assistant_token_id_ < 0) { + // Try with space prefix (GPT-2 style) + assistant_token_id_ = find_token("Ġassistant"); + } + + // Newline token + newline_token_id_ = find_token("Ċ"); // GPT-2 encoding for '\n' + if (newline_token_id_ < 0) { + newline_token_id_ = find_token("\n"); + } + + loaded_ = true; + return true; +} + +std::pair TextTokenizer::get_min_pair( + const std::vector & word) const { + + std::pair min_pair; + int32_t min_rank = std::numeric_limits::max(); + + for (size_t i = 0; i + 1 < word.size(); i++) { + auto pair = std::make_pair(word[i], word[i + 1]); + auto it = bpe_ranks_.find(pair); + if (it != bpe_ranks_.end() && it->second < min_rank) { + min_rank = it->second; + min_pair = pair; + } + } + + return min_pair; +} + +std::vector TextTokenizer::bpe(const std::string & token) const { + if (token.empty()) { + return {}; + } + + // Split into unicode characters + std::vector word; + size_t i = 0; + while (i < token.size()) { + size_t len = utf8_len(token[i]); + word.push_back(token.substr(i, len)); + i += len; + } + + if (word.size() == 1) { + return word; + } + + // Iteratively merge pairs + while (true) { + auto min_pair = get_min_pair(word); + if (min_pair.first.empty()) { + break; // No more merges possible + } + + // Merge all occurrences of the pair + std::vector new_word; + size_t j = 0; + while (j < word.size()) { + if (j + 1 < word.size() && + word[j] == min_pair.first && + word[j + 1] == min_pair.second) { + new_word.push_back(min_pair.first + min_pair.second); + j += 2; + } else { + new_word.push_back(word[j]); + j += 1; + } + } + word = std::move(new_word); + + if (word.size() == 1) { + break; + } + } + + return word; +} + +std::vector TextTokenizer::encode(const std::string & text) const { + if (!loaded_) { + return {}; + } + + std::vector tokens; + + // Convert text to GPT-2 unicode representation + std::string unicode_text = bytes_to_unicode(text); + + // Simple word splitting (no regex pre-tokenization for now) + // Split on spaces but keep the space with the following word (GPT-2 style) + std::vector words; + std::string current_word; + + size_t i = 0; + while (i < unicode_text.size()) { + size_t len = utf8_len(unicode_text[i]); + std::string ch = unicode_text.substr(i, len); + + // Check if this is a space (Ġ in GPT-2 encoding) + if (ch == "Ġ") { + if (!current_word.empty()) { + words.push_back(current_word); + current_word.clear(); + } + current_word = ch; // Start new word with space + } else { + current_word += ch; + } + i += len; + } + if (!current_word.empty()) { + words.push_back(current_word); + } + + // BPE encode each word + for (const auto & word : words) { + auto bpe_tokens = bpe(word); + for (const auto & tok : bpe_tokens) { + auto it = vocab_.find(tok); + if (it != vocab_.end()) { + tokens.push_back(it->second); + } else { + // Unknown token - encode as bytes + for (unsigned char c : tok) { + std::string byte_tok = BYTE_TO_UNICODE[c]; + auto byte_it = vocab_.find(byte_tok); + if (byte_it != vocab_.end()) { + tokens.push_back(byte_it->second); + } + } + } + } + } + + return tokens; +} + +std::vector TextTokenizer::encode_for_tts(const std::string & text) const { + if (!loaded_) { + return {}; + } + + // Format: <|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n + std::vector tokens; + + // <|im_start|> + tokens.push_back(config_.bos_token_id); + + // assistant + tokens.push_back(assistant_token_id_); + + // \n + tokens.push_back(newline_token_id_); + + // Encode the text + auto text_tokens = encode(text); + tokens.insert(tokens.end(), text_tokens.begin(), text_tokens.end()); + + // <|im_end|> + tokens.push_back(config_.eos_token_id); + + // \n + tokens.push_back(newline_token_id_); + + // <|im_start|> + tokens.push_back(config_.bos_token_id); + + // assistant + tokens.push_back(assistant_token_id_); + + // \n + tokens.push_back(newline_token_id_); + + return tokens; +} + +std::string TextTokenizer::decode(const std::vector & tokens) const { + std::string result; + for (int32_t token : tokens) { + result += decode_token(token); + } + return result; +} + +std::string TextTokenizer::decode_token(int32_t token_id) const { + if (token_id < 0 || token_id >= (int32_t)id_to_token_.size()) { + return ""; + } + + const std::string & token = id_to_token_[token_id]; + + // Convert from GPT-2 unicode back to bytes + return unicode_to_bytes(token); +} + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/text_tokenizer.h b/otherarch/qwen3tts/text_tokenizer.h new file mode 100644 index 000000000..4944f4966 --- /dev/null +++ b/otherarch/qwen3tts/text_tokenizer.h @@ -0,0 +1,88 @@ +#pragma once + +#include "gguf.h" + +#include +#include +#include +#include + +namespace qwen3_tts { + +// BPE tokenizer configuration +struct tokenizer_config { + int32_t vocab_size = 151936; + int32_t pad_token_id = 151643; + int32_t eos_token_id = 151645; // <|im_end|> + int32_t bos_token_id = 151644; // <|im_start|> +}; + +// Text tokenizer class (BPE-based, GPT-2 style with Qwen2 pre-tokenization) +class TextTokenizer { +public: + TextTokenizer(); + ~TextTokenizer(); + + // Load tokenizer from GGUF file + bool load_from_gguf(struct gguf_context * ctx); + + // Encode text to token IDs + std::vector encode(const std::string & text) const; + + // Encode with TTS format: <|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n + std::vector encode_for_tts(const std::string & text) const; + + // Decode token IDs to text + std::string decode(const std::vector & tokens) const; + + // Decode single token + std::string decode_token(int32_t token_id) const; + + // Get configuration + const tokenizer_config & get_config() const { return config_; } + + // Get error message + const std::string & get_error() const { return error_msg_; } + + // Check if loaded + bool is_loaded() const { return loaded_; } + + // Get special token IDs + int32_t bos_token_id() const { return config_.bos_token_id; } + int32_t eos_token_id() const { return config_.eos_token_id; } + int32_t pad_token_id() const { return config_.pad_token_id; } + +private: + tokenizer_config config_; + std::string error_msg_; + bool loaded_ = false; + + // Vocabulary: token string -> token ID + std::unordered_map vocab_; + + // Reverse vocabulary: token ID -> token string + std::vector id_to_token_; + + // BPE merges: pair -> rank (lower rank = higher priority) + std::map, int32_t> bpe_ranks_; + + // Special token for "assistant" and newline + int32_t assistant_token_id_ = 77091; + int32_t newline_token_id_ = 198; // '\n' encoded + + // Helper: convert bytes to unicode (GPT-2 style byte encoding) + static std::string bytes_to_unicode(const std::string & text); + static std::string unicode_to_bytes(const std::string & text); + + // Helper: get UTF-8 character length + static size_t utf8_len(char c); + + // BPE encoding for a single word + std::vector bpe(const std::string & token) const; + + // Find the pair with lowest rank in a sequence + std::pair get_min_pair( + const std::vector & word) const; +}; + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/tts_transformer.cpp b/otherarch/qwen3tts/tts_transformer.cpp new file mode 100644 index 000000000..cff2b5423 --- /dev/null +++ b/otherarch/qwen3tts/tts_transformer.cpp @@ -0,0 +1,2899 @@ +#include "tts_transformer.h" +#include "gguf_loader.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace qwen3_tts { + +TTSTransformer::TTSTransformer() = default; + +TTSTransformer::~TTSTransformer() { + unload_model(); +} + +void TTSTransformer::unload_model() { + free_tts_kv_cache(state_.cache); + free_tts_kv_cache(state_.code_pred_cache); + free_transformer_model(model_); + + coreml_code_predictor_.unload(); + use_coreml_code_predictor_ = false; + coreml_code_predictor_path_.clear(); + skip_ggml_code_pred_layers_ = false; + + if (state_.sched) { + ggml_backend_sched_free(state_.sched); + state_.sched = nullptr; + } + if (state_.backend) { + release_preferred_backend(state_.backend); + state_.backend = nullptr; + } + if (state_.backend_cpu) { + ggml_backend_free(state_.backend_cpu); + state_.backend_cpu = nullptr; + } + + state_.compute_meta.clear(); + last_hidden_.clear(); + embd_row_fp16_scratch_.clear(); +} + +bool TTSTransformer::load_model(const std::string & model_path) { + unload_model(); + + skip_ggml_code_pred_layers_ = false; +#if defined(__APPLE__) + const char * use_coreml_env = std::getenv("QWEN3_TTS_USE_COREML"); + bool coreml_disabled = false; + if (use_coreml_env && use_coreml_env[0] != '\0') { + std::string use_coreml = use_coreml_env; + std::transform(use_coreml.begin(), use_coreml.end(), use_coreml.begin(), + [](unsigned char c) { return (char) std::tolower(c); }); + coreml_disabled = use_coreml == "0" || use_coreml == "false" || + use_coreml == "off" || use_coreml == "no"; + } + + if (!coreml_disabled) { + std::string coreml_path; + const char * override_env = std::getenv("QWEN3_TTS_COREML_MODEL"); + if (override_env && override_env[0] != '\0') { + coreml_path = override_env; + } else { + size_t slash = model_path.find_last_of("/\\"); + const std::string model_dir = (slash == std::string::npos) ? "." : model_path.substr(0, slash); + coreml_path = model_dir + "/coreml/code_predictor.mlpackage"; + } + + struct stat st = {}; + if (stat(coreml_path.c_str(), &st) == 0) { + // Skip GGML code-predictor weights when CoreML package is present. + skip_ggml_code_pred_layers_ = true; + } else if (use_coreml_env && use_coreml_env[0] != '\0') { + // Explicit opt-in should remain strict to surface configuration errors. + skip_ggml_code_pred_layers_ = true; + } + } +#endif + + struct ggml_context * meta_ctx = nullptr; + struct gguf_init_params params = { + /*.no_alloc =*/ true, + /*.ctx =*/ &meta_ctx, + }; + + struct gguf_context * ctx = gguf_init_from_file(model_path.c_str(), params); + if (!ctx) { + error_msg_ = "Failed to open GGUF file: " + model_path; + return false; + } + + if (!parse_config(ctx)) { + gguf_free(ctx); + if (meta_ctx) ggml_free(meta_ctx); + return false; + } + + if (!create_tensors(ctx)) { + gguf_free(ctx); + if (meta_ctx) ggml_free(meta_ctx); + return false; + } + + if (!load_tensor_data(model_path, ctx)) { + free_transformer_model(model_); + gguf_free(ctx); + if (meta_ctx) ggml_free(meta_ctx); + return false; + } + + gguf_free(ctx); + if (meta_ctx) ggml_free(meta_ctx); + + state_.backend = init_preferred_backend("TTSTransformer", &error_msg_); + if (!state_.backend) { + return false; + } + ggml_backend_dev_t device = ggml_backend_get_device(state_.backend); + const char * device_name = device ? ggml_backend_dev_name(device) : "Unknown"; + fprintf(stderr, " TTSTransformer backend: %s\n", device_name); + + if (device && ggml_backend_dev_type(device) != GGML_BACKEND_DEVICE_TYPE_CPU) { + state_.backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + if (!state_.backend_cpu) { + error_msg_ = "Failed to initialize CPU fallback backend for TTSTransformer"; + return false; + } + } + + std::vector backends; + backends.push_back(state_.backend); + if (state_.backend_cpu) { + backends.push_back(state_.backend_cpu); + } + state_.sched = ggml_backend_sched_new(backends.data(), nullptr, (int)backends.size(), QWEN3_TTS_MAX_NODES, false, true); + if (!state_.sched) { + error_msg_ = "Failed to create backend scheduler"; + return false; + } + + state_.compute_meta.resize(ggml_tensor_overhead() * QWEN3_TTS_MAX_NODES + ggml_graph_overhead()); + + if (!try_init_coreml_code_predictor(model_path)) { + return false; + } + + return true; +} + +bool TTSTransformer::try_init_coreml_code_predictor(const std::string & model_path) { + use_coreml_code_predictor_ = false; + coreml_code_predictor_path_.clear(); + + const char * use_coreml_env = std::getenv("QWEN3_TTS_USE_COREML"); + bool coreml_disabled = false; + if (use_coreml_env && use_coreml_env[0] != '\0') { + std::string use_coreml = use_coreml_env; + std::transform(use_coreml.begin(), use_coreml.end(), use_coreml.begin(), + [](unsigned char c) { return (char) std::tolower(c); }); + coreml_disabled = use_coreml == "0" || use_coreml == "false" || + use_coreml == "off" || use_coreml == "no"; + } + + if (coreml_disabled) { + return true; + } + +#if !defined(__APPLE__) + if (use_coreml_env && use_coreml_env[0] != '\0') { + fprintf(stderr, " CoreML code predictor requested but this build is not on Apple platform\n"); + } + return true; +#else + std::string coreml_path; + const char * override_env = std::getenv("QWEN3_TTS_COREML_MODEL"); + if (override_env && override_env[0] != '\0') { + coreml_path = override_env; + } else { + size_t slash = model_path.find_last_of("/\\"); + const std::string model_dir = (slash == std::string::npos) ? "." : model_path.substr(0, slash); + coreml_path = model_dir + "/coreml/code_predictor.mlpackage"; + } + + if (!coreml_code_predictor_.load(coreml_path, model_.config.n_codebooks - 1)) { + if (skip_ggml_code_pred_layers_) { + error_msg_ = "CoreML code predictor load failed in strict mode: " + coreml_code_predictor_.get_error(); + return false; + } else { + fprintf(stderr, " CoreML code predictor load failed: %s\n", + coreml_code_predictor_.get_error().c_str()); + fprintf(stderr, " Falling back to GGML code predictor\n"); + return true; + } + } + + use_coreml_code_predictor_ = true; + coreml_code_predictor_path_ = coreml_path; + fprintf(stderr, " CoreML code predictor enabled: %s\n", coreml_code_predictor_path_.c_str()); + return true; +#endif +} + +bool TTSTransformer::parse_config(struct gguf_context * ctx) { + auto get_u32_any = [&](std::initializer_list keys, int32_t default_val) -> int32_t { + for (const char * key : keys) { + int64_t idx = gguf_find_key(ctx, key); + if (idx >= 0) { + return (int32_t)gguf_get_val_u32(ctx, idx); + } + } + return default_val; + }; + + auto get_f32_any = [&](std::initializer_list keys, float default_val) -> float { + for (const char * key : keys) { + int64_t idx = gguf_find_key(ctx, key); + if (idx >= 0) { + return gguf_get_val_f32(ctx, idx); + } + } + return default_val; + }; + + auto & cfg = model_.config; + cfg.text_vocab_size = get_u32_any({ + "qwen3-tts.text.vocab_size", + "qwen3-tts.text_vocab_size", + }, 151936); + cfg.text_embd_dim = get_u32_any({ + "qwen3-tts.text.embedding_dim", + "qwen3-tts.text_hidden_size", + }, 2048); + cfg.hidden_size = get_u32_any({ + "qwen3-tts.talker.embedding_length", + "qwen3-tts.embedding_length", + }, 1024); + cfg.n_layers = get_u32_any({ + "qwen3-tts.talker.block_count", + "qwen3-tts.block_count", + }, 28); + cfg.n_attention_heads = get_u32_any({ + "qwen3-tts.talker.attention.head_count", + "qwen3-tts.attention.head_count", + }, 16); + cfg.n_key_value_heads = get_u32_any({ + "qwen3-tts.talker.attention.head_count_kv", + "qwen3-tts.attention.head_count_kv", + }, 8); + cfg.intermediate_size = get_u32_any({ + "qwen3-tts.talker.feed_forward_length", + "qwen3-tts.feed_forward_length", + }, 3072); + cfg.head_dim = get_u32_any({ + "qwen3-tts.talker.attention.key_length", + "qwen3-tts.attention.key_length", + }, 128); + cfg.rms_norm_eps = get_f32_any({ + "qwen3-tts.talker.attention.layer_norm_rms_epsilon", + "qwen3-tts.attention.layer_norm_rms_epsilon", + }, 1e-6f); + cfg.rope_theta = get_f32_any({ + "qwen3-tts.talker.rope.freq_base", + "qwen3-tts.rope.freq_base", + }, 1000000.0f); + + cfg.codec_vocab_size = get_u32_any({ + "qwen3-tts.talker.codec_vocab_size", + "qwen3-tts.vocab_size", + }, 3072); + cfg.n_codebooks = get_u32_any({ + "qwen3-tts.talker.num_codebooks", + "qwen3-tts.num_code_groups", + }, 16); + + cfg.code_pred_layers = get_u32_any({ + "qwen3-tts.code_pred.layer_count", + "qwen3-tts.code_predictor.layer_count", + }, 5); + cfg.code_pred_vocab_size = get_u32_any({ + "qwen3-tts.code_pred.vocab_size", + "qwen3-tts.code_predictor.vocab_size", + }, 2048); + + cfg.codec_pad_id = get_u32_any({ + "qwen3-tts.codec.pad_id", + }, 2148); + cfg.codec_bos_id = get_u32_any({ + "qwen3-tts.codec.bos_id", + }, 2149); + cfg.codec_eos_id = get_u32_any({ + "qwen3-tts.codec.eos_id", + "qwen3-tts.codec.eos_token_id", + }, 2150); + + cfg.tts_bos_token_id = get_u32_any({ + "qwen3-tts.tts_bos_token_id", + "qwen3-tts.tts.bos_token_id", + "qwen3-tts.tts.bos_id", + }, 151672); + cfg.tts_eos_token_id = get_u32_any({ + "qwen3-tts.tts_eos_token_id", + "qwen3-tts.tts.eos_token_id", + "qwen3-tts.tts.eos_id", + }, 151673); + cfg.tts_pad_token_id = get_u32_any({ + "qwen3-tts.tts_pad_token_id", + "qwen3-tts.tts.pad_token_id", + "qwen3-tts.tts.pad_id", + }, 151671); + + cfg.codec_think_id = get_u32_any({ + "qwen3-tts.codec.think_id", + "qwen3-tts.codec_think_id", + }, 2154); + cfg.codec_nothink_id = get_u32_any({ + "qwen3-tts.codec.nothink_id", + "qwen3-tts.codec_nothink_id", + }, 2155); + cfg.codec_think_bos_id = get_u32_any({ + "qwen3-tts.codec.think_bos_id", + "qwen3-tts.codec_think_bos_id", + }, 2156); + cfg.codec_think_eos_id = get_u32_any({ + "qwen3-tts.codec.think_eos_id", + "qwen3-tts.codec_think_eos_id", + }, 2157); + + cfg.english_language_id = get_u32_any({ + "qwen3-tts.language.english_id", + "qwen3-tts.codec.language.english_id", + "qwen3-tts.language_id", + }, 2050); + + return true; +} + +bool TTSTransformer::create_tensors(struct gguf_context * ctx) { + const int64_t n_tensors = gguf_get_n_tensors(ctx); + const auto & cfg = model_.config; + + const size_t ctx_size = n_tensors * ggml_tensor_overhead(); + struct ggml_init_params params = { + /*.mem_size =*/ ctx_size, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + model_.ctx = ggml_init(params); + if (!model_.ctx) { + error_msg_ = "Failed to create GGML context"; + return false; + } + + model_.layers.resize(cfg.n_layers); + model_.code_pred_layers.resize(cfg.code_pred_layers); + model_.code_pred_embd.resize(cfg.n_codebooks - 1); + model_.code_pred_head.resize(cfg.n_codebooks - 1); + + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = gguf_get_tensor_name(ctx, i); + enum ggml_type type = gguf_get_tensor_type(ctx, i); + + int64_t ne[GGML_MAX_DIMS] = {1, 1, 1, 1}; + int n_dims = 0; + + if (strstr(name, "spk_enc.") || strstr(name, "tok_")) { + continue; + } + + if (strstr(name, "talker.text_embd.weight")) { + ne[0] = cfg.text_embd_dim; + ne[1] = cfg.text_vocab_size; + n_dims = 2; + } else if (strstr(name, "talker.text_proj.fc1.weight")) { + ne[0] = cfg.text_embd_dim; + ne[1] = cfg.text_embd_dim; + n_dims = 2; + } else if (strstr(name, "talker.text_proj.fc1.bias")) { + ne[0] = cfg.text_embd_dim; + n_dims = 1; + } else if (strstr(name, "talker.text_proj.fc2.weight")) { + ne[0] = cfg.text_embd_dim; + ne[1] = cfg.hidden_size; + n_dims = 2; + } else if (strstr(name, "talker.text_proj.fc2.bias")) { + ne[0] = cfg.hidden_size; + n_dims = 1; + } else if (strstr(name, "talker.codec_embd.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.codec_vocab_size; + n_dims = 2; + } else if (strstr(name, "talker.codec_head.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.codec_vocab_size; + n_dims = 2; + } else if (strstr(name, "talker.output_norm.weight")) { + ne[0] = cfg.hidden_size; + n_dims = 1; + } else if (strstr(name, "talker.blk.")) { + int layer_idx = -1; + if (sscanf(name, "talker.blk.%d.", &layer_idx) == 1 && + layer_idx >= 0 && layer_idx < cfg.n_layers) { + + if (strstr(name, "attn_norm.weight")) { + ne[0] = cfg.hidden_size; + n_dims = 1; + } else if (strstr(name, "attn_q_norm.weight")) { + ne[0] = cfg.head_dim; + n_dims = 1; + } else if (strstr(name, "attn_k_norm.weight")) { + ne[0] = cfg.head_dim; + n_dims = 1; + } else if (strstr(name, "attn_q.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.n_attention_heads * cfg.head_dim; + n_dims = 2; + } else if (strstr(name, "attn_k.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.n_key_value_heads * cfg.head_dim; + n_dims = 2; + } else if (strstr(name, "attn_v.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.n_key_value_heads * cfg.head_dim; + n_dims = 2; + } else if (strstr(name, "attn_output.weight")) { + ne[0] = cfg.n_attention_heads * cfg.head_dim; + ne[1] = cfg.hidden_size; + n_dims = 2; + } else if (strstr(name, "ffn_norm.weight")) { + ne[0] = cfg.hidden_size; + n_dims = 1; + } else if (strstr(name, "ffn_gate.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.intermediate_size; + n_dims = 2; + } else if (strstr(name, "ffn_up.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.intermediate_size; + n_dims = 2; + } else if (strstr(name, "ffn_down.weight")) { + ne[0] = cfg.intermediate_size; + ne[1] = cfg.hidden_size; + n_dims = 2; + } else { + continue; + } + } else { + continue; + } + } else if (strstr(name, "code_pred.blk.")) { + if (skip_ggml_code_pred_layers_) { + continue; + } + int layer_idx = -1; + if (sscanf(name, "code_pred.blk.%d.", &layer_idx) == 1 && + layer_idx >= 0 && layer_idx < cfg.code_pred_layers) { + + if (strstr(name, "attn_norm.weight")) { + ne[0] = cfg.hidden_size; + n_dims = 1; + } else if (strstr(name, "attn_q_norm.weight")) { + ne[0] = cfg.head_dim; + n_dims = 1; + } else if (strstr(name, "attn_k_norm.weight")) { + ne[0] = cfg.head_dim; + n_dims = 1; + } else if (strstr(name, "attn_q.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.n_attention_heads * cfg.head_dim; + n_dims = 2; + } else if (strstr(name, "attn_k.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.n_key_value_heads * cfg.head_dim; + n_dims = 2; + } else if (strstr(name, "attn_v.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.n_key_value_heads * cfg.head_dim; + n_dims = 2; + } else if (strstr(name, "attn_output.weight")) { + ne[0] = cfg.n_attention_heads * cfg.head_dim; + ne[1] = cfg.hidden_size; + n_dims = 2; + } else if (strstr(name, "ffn_norm.weight")) { + ne[0] = cfg.hidden_size; + n_dims = 1; + } else if (strstr(name, "ffn_gate.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.intermediate_size; + n_dims = 2; + } else if (strstr(name, "ffn_up.weight")) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.intermediate_size; + n_dims = 2; + } else if (strstr(name, "ffn_down.weight")) { + ne[0] = cfg.intermediate_size; + ne[1] = cfg.hidden_size; + n_dims = 2; + } else { + continue; + } + } else { + continue; + } + } else if (strstr(name, "code_pred.codec_embd.")) { + int cb_idx = -1; + if (sscanf(name, "code_pred.codec_embd.%d.weight", &cb_idx) == 1 && + cb_idx >= 0 && cb_idx < cfg.n_codebooks - 1) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.code_pred_vocab_size; + n_dims = 2; + } else { + continue; + } + } else if (strstr(name, "code_pred.lm_head.")) { + if (skip_ggml_code_pred_layers_) { + continue; + } + int cb_idx = -1; + if (sscanf(name, "code_pred.lm_head.%d.weight", &cb_idx) == 1 && + cb_idx >= 0 && cb_idx < cfg.n_codebooks - 1) { + ne[0] = cfg.hidden_size; + ne[1] = cfg.code_pred_vocab_size; + n_dims = 2; + } else { + continue; + } + } else if (strstr(name, "code_pred.output_norm.weight")) { + if (skip_ggml_code_pred_layers_) { + continue; + } + ne[0] = cfg.hidden_size; + n_dims = 1; + } else { + continue; + } + + struct ggml_tensor * tensor = ggml_new_tensor(model_.ctx, type, n_dims, ne); + if (!tensor) { + error_msg_ = "Failed to create tensor: " + std::string(name); + return false; + } + ggml_set_name(tensor, name); + model_.tensors[name] = tensor; + + if (strstr(name, "talker.text_embd.weight")) { + model_.text_embd = tensor; + } else if (strstr(name, "talker.text_proj.fc1.weight")) { + model_.text_proj_fc1 = tensor; + } else if (strstr(name, "talker.text_proj.fc1.bias")) { + model_.text_proj_fc1_bias = tensor; + } else if (strstr(name, "talker.text_proj.fc2.weight")) { + model_.text_proj_fc2 = tensor; + } else if (strstr(name, "talker.text_proj.fc2.bias")) { + model_.text_proj_fc2_bias = tensor; + } else if (strstr(name, "talker.codec_embd.weight")) { + model_.codec_embd = tensor; + } else if (strstr(name, "talker.codec_head.weight")) { + model_.codec_head = tensor; + } else if (strstr(name, "talker.output_norm.weight")) { + model_.output_norm = tensor; + } else if (strstr(name, "talker.blk.")) { + int layer_idx = -1; + sscanf(name, "talker.blk.%d.", &layer_idx); + if (layer_idx >= 0 && layer_idx < cfg.n_layers) { + auto & layer = model_.layers[layer_idx]; + if (strstr(name, "attn_norm.weight")) layer.attn_norm = tensor; + else if (strstr(name, "attn_q_norm.weight")) layer.attn_q_norm = tensor; + else if (strstr(name, "attn_k_norm.weight")) layer.attn_k_norm = tensor; + else if (strstr(name, "attn_q.weight")) layer.attn_q = tensor; + else if (strstr(name, "attn_k.weight")) layer.attn_k = tensor; + else if (strstr(name, "attn_v.weight")) layer.attn_v = tensor; + else if (strstr(name, "attn_output.weight")) layer.attn_output = tensor; + else if (strstr(name, "ffn_norm.weight")) layer.ffn_norm = tensor; + else if (strstr(name, "ffn_gate.weight")) layer.ffn_gate = tensor; + else if (strstr(name, "ffn_up.weight")) layer.ffn_up = tensor; + else if (strstr(name, "ffn_down.weight")) layer.ffn_down = tensor; + } + } else if (strstr(name, "code_pred.blk.")) { + int layer_idx = -1; + sscanf(name, "code_pred.blk.%d.", &layer_idx); + if (layer_idx >= 0 && layer_idx < cfg.code_pred_layers) { + auto & layer = model_.code_pred_layers[layer_idx]; + if (strstr(name, "attn_norm.weight")) layer.attn_norm = tensor; + else if (strstr(name, "attn_q_norm.weight")) layer.attn_q_norm = tensor; + else if (strstr(name, "attn_k_norm.weight")) layer.attn_k_norm = tensor; + else if (strstr(name, "attn_q.weight")) layer.attn_q = tensor; + else if (strstr(name, "attn_k.weight")) layer.attn_k = tensor; + else if (strstr(name, "attn_v.weight")) layer.attn_v = tensor; + else if (strstr(name, "attn_output.weight")) layer.attn_output = tensor; + else if (strstr(name, "ffn_norm.weight")) layer.ffn_norm = tensor; + else if (strstr(name, "ffn_gate.weight")) layer.ffn_gate = tensor; + else if (strstr(name, "ffn_up.weight")) layer.ffn_up = tensor; + else if (strstr(name, "ffn_down.weight")) layer.ffn_down = tensor; + } + } else if (strstr(name, "code_pred.codec_embd.")) { + int cb_idx = -1; + sscanf(name, "code_pred.codec_embd.%d.weight", &cb_idx); + if (cb_idx >= 0 && cb_idx < cfg.n_codebooks - 1) { + model_.code_pred_embd[cb_idx] = tensor; + } + } else if (strstr(name, "code_pred.lm_head.")) { + int cb_idx = -1; + sscanf(name, "code_pred.lm_head.%d.weight", &cb_idx); + if (cb_idx >= 0 && cb_idx < cfg.n_codebooks - 1) { + model_.code_pred_head[cb_idx] = tensor; + } + } else if (strstr(name, "code_pred.output_norm.weight")) { + model_.code_pred_output_norm = tensor; + } + } + + return true; + } + +bool TTSTransformer::load_tensor_data(const std::string & path, struct gguf_context * ctx) { + ggml_backend_t backend = init_preferred_backend("TTSTransformer", &error_msg_); + if (!backend) { + return false; + } + + model_.buffer = ggml_backend_alloc_ctx_tensors(model_.ctx, backend); + if (!model_.buffer) { + error_msg_ = "Failed to allocate tensor buffer"; + release_preferred_backend(backend); + return false; + } + + FILE * f = fopen(path.c_str(), "rb"); + if (!f) { + error_msg_ = "Failed to open file for reading: " + path; + release_preferred_backend(backend); + return false; + } + + const size_t data_offset = gguf_get_data_offset(ctx); + const int64_t n_tensors = gguf_get_n_tensors(ctx); + std::vector read_buf; + + for (int64_t i = 0; i < n_tensors; ++i) { + const char * name = gguf_get_tensor_name(ctx, i); + size_t offset = gguf_get_tensor_offset(ctx, i); + + auto it = model_.tensors.find(name); + if (it == model_.tensors.end()) { + continue; + } + + struct ggml_tensor * tensor = it->second; + size_t nbytes = ggml_nbytes(tensor); + + read_buf.resize(nbytes); + + if (fseek(f, data_offset + offset, SEEK_SET) != 0) { + error_msg_ = "Failed to seek to tensor data: " + std::string(name); + fclose(f); + release_preferred_backend(backend); + return false; + } + + if (fread(read_buf.data(), 1, nbytes, f) != nbytes) { + error_msg_ = "Failed to read tensor data: " + std::string(name); + fclose(f); + release_preferred_backend(backend); + return false; + } + + ggml_backend_tensor_set(tensor, read_buf.data(), 0, nbytes); + } + + fclose(f); + release_preferred_backend(backend); + + return true; +} + +bool TTSTransformer::init_kv_cache(int32_t n_ctx) { + const auto & cfg = model_.config; + + free_tts_kv_cache(state_.cache); + + state_.cache.n_ctx = n_ctx; + state_.cache.n_used = 0; + state_.cache.head_dim = cfg.head_dim; + state_.cache.n_kv_heads = cfg.n_key_value_heads; + state_.cache.n_layers = cfg.n_layers; + + const size_t n_tensors = cfg.n_layers * 2; + const size_t ctx_size = n_tensors * ggml_tensor_overhead(); + + struct ggml_init_params params = { + /*.mem_size =*/ ctx_size, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + state_.cache.ctx = ggml_init(params); + if (!state_.cache.ctx) { + error_msg_ = "Failed to create KV cache context"; + return false; + } + + state_.cache.k_cache.resize(cfg.n_layers); + state_.cache.v_cache.resize(cfg.n_layers); + + for (int il = 0; il < cfg.n_layers; ++il) { + state_.cache.k_cache[il] = ggml_new_tensor_3d( + state_.cache.ctx, GGML_TYPE_F16, + cfg.head_dim, cfg.n_key_value_heads, n_ctx); + ggml_format_name(state_.cache.k_cache[il], "k_cache_%d", il); + + state_.cache.v_cache[il] = ggml_new_tensor_3d( + state_.cache.ctx, GGML_TYPE_F16, + cfg.head_dim, cfg.n_key_value_heads, n_ctx); + ggml_format_name(state_.cache.v_cache[il], "v_cache_%d", il); + } + + state_.cache.buffer = ggml_backend_alloc_ctx_tensors(state_.cache.ctx, state_.backend); + if (!state_.cache.buffer) { + error_msg_ = "Failed to allocate KV cache buffer"; + return false; + } + + return true; +} + +void TTSTransformer::clear_kv_cache() { + state_.cache.n_used = 0; +} + +bool TTSTransformer::init_code_pred_kv_cache(int32_t n_ctx) { + const auto & cfg = model_.config; + + free_tts_kv_cache(state_.code_pred_cache); + + state_.code_pred_cache.n_ctx = n_ctx; + state_.code_pred_cache.n_used = 0; + state_.code_pred_cache.head_dim = cfg.head_dim; + state_.code_pred_cache.n_kv_heads = cfg.n_key_value_heads; + state_.code_pred_cache.n_layers = cfg.code_pred_layers; + + const size_t n_tensors = cfg.code_pred_layers * 2; + const size_t ctx_size = n_tensors * ggml_tensor_overhead(); + + struct ggml_init_params params = { + /*.mem_size =*/ ctx_size, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + state_.code_pred_cache.ctx = ggml_init(params); + if (!state_.code_pred_cache.ctx) { + error_msg_ = "Failed to create code predictor KV cache context"; + return false; + } + + state_.code_pred_cache.k_cache.resize(cfg.code_pred_layers); + state_.code_pred_cache.v_cache.resize(cfg.code_pred_layers); + + for (int il = 0; il < cfg.code_pred_layers; ++il) { + state_.code_pred_cache.k_cache[il] = ggml_new_tensor_3d( + state_.code_pred_cache.ctx, GGML_TYPE_F16, + cfg.head_dim, cfg.n_key_value_heads, n_ctx); + ggml_format_name(state_.code_pred_cache.k_cache[il], "code_pred_k_cache_%d", il); + + state_.code_pred_cache.v_cache[il] = ggml_new_tensor_3d( + state_.code_pred_cache.ctx, GGML_TYPE_F16, + cfg.head_dim, cfg.n_key_value_heads, n_ctx); + ggml_format_name(state_.code_pred_cache.v_cache[il], "code_pred_v_cache_%d", il); + } + + state_.code_pred_cache.buffer = ggml_backend_alloc_ctx_tensors(state_.code_pred_cache.ctx, state_.backend); + if (!state_.code_pred_cache.buffer) { + error_msg_ = "Failed to allocate code predictor KV cache buffer"; + return false; + } + + return true; +} + +void TTSTransformer::clear_code_pred_kv_cache() { + state_.code_pred_cache.n_used = 0; +} + +bool TTSTransformer::lookup_embedding_rows(struct ggml_tensor * embedding, const int32_t * token_ids, + int32_t n_tokens, const char * input_name, + const char * output_name, std::vector & output) { + if (!model_.ctx) { + error_msg_ = "Model not loaded"; + return false; + } + if (!embedding) { + error_msg_ = "Embedding tensor not found"; + return false; + } + if (n_tokens <= 0) { + output.clear(); + return true; + } + + const int32_t embd_dim = (int32_t) embedding->ne[0]; + if (n_tokens <= 32 && + (embedding->type == GGML_TYPE_F16 || embedding->type == GGML_TYPE_F32)) { + output.resize((size_t) embd_dim * n_tokens); + for (int32_t t = 0; t < n_tokens; ++t) { + if (!lookup_single_embedding_row(embedding, token_ids[t], + output.data() + (size_t) t * embd_dim)) { + return false; + } + } + return true; + } + + struct ggml_init_params params = { + /*.mem_size =*/ state_.compute_meta.size(), + /*.mem_buffer =*/ state_.compute_meta.data(), + /*.no_alloc =*/ true, + }; + + struct ggml_context * ctx0 = ggml_init(params); + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false); + + struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_name(inp_tokens, input_name); + ggml_set_input(inp_tokens); + + struct ggml_tensor * rows = ggml_get_rows(ctx0, embedding, inp_tokens); + rows = ggml_cast(ctx0, rows, GGML_TYPE_F32); + ggml_set_name(rows, output_name); + ggml_set_output(rows); + + ggml_build_forward_expand(gf, rows); + + if (!ggml_backend_sched_alloc_graph(state_.sched, gf)) { + error_msg_ = "Failed to allocate embedding lookup graph"; + ggml_free(ctx0); + return false; + } + + struct ggml_tensor * inp = ggml_graph_get_tensor(gf, input_name); + ggml_backend_tensor_set(inp, token_ids, 0, n_tokens * sizeof(int32_t)); + + if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) { + error_msg_ = "Failed to compute embedding lookup graph"; + ggml_backend_sched_reset(state_.sched); + ggml_free(ctx0); + return false; + } + + struct ggml_tensor * out = ggml_graph_get_tensor(gf, output_name); + if (!out) { + error_msg_ = "Failed to find embedding lookup output tensor"; + ggml_backend_sched_reset(state_.sched); + ggml_free(ctx0); + return false; + } + + output.resize((size_t)embedding->ne[0] * n_tokens); + ggml_backend_tensor_get(out, output.data(), 0, output.size() * sizeof(float)); + + ggml_backend_sched_reset(state_.sched); + ggml_free(ctx0); + return true; +} + +bool TTSTransformer::lookup_single_embedding_row(struct ggml_tensor * embedding, int32_t token_id, + float * out_row) { + if (!embedding) { + error_msg_ = "Embedding tensor not found"; + return false; + } + if (!out_row) { + error_msg_ = "Embedding output row is null"; + return false; + } + + const int64_t embd_dim = embedding->ne[0]; + const int64_t vocab_size = embedding->ne[1]; + if (token_id < 0 || token_id >= vocab_size) { + error_msg_ = "Embedding token ID out of range"; + return false; + } + + const size_t row_offset = (size_t) token_id * embedding->nb[1]; + if (embedding->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(embedding, out_row, row_offset, (size_t) embd_dim * sizeof(float)); + return true; + } + if (embedding->type == GGML_TYPE_F16) { + embd_row_fp16_scratch_.resize((size_t) embd_dim); + ggml_backend_tensor_get(embedding, embd_row_fp16_scratch_.data(), + row_offset, (size_t) embd_dim * sizeof(ggml_fp16_t)); + for (int64_t i = 0; i < embd_dim; ++i) { + out_row[i] = ggml_fp16_to_fp32(embd_row_fp16_scratch_[i]); + } + return true; + } + + std::vector single_token = { token_id }; + std::vector single_out; + if (!lookup_embedding_rows(embedding, single_token.data(), 1, + "inp_compat_embed", "out_compat_embed", single_out)) { + return false; + } + memcpy(out_row, single_out.data(), (size_t) embd_dim * sizeof(float)); + return true; +} + +bool TTSTransformer::project_text_tokens(const int32_t * text_tokens, int32_t n_tokens, + std::vector & output) { + if (!model_.ctx) { + error_msg_ = "Model not loaded"; + return false; + } + if (n_tokens <= 0) { + output.clear(); + return true; + } + + struct ggml_init_params params = { + /*.mem_size =*/ state_.compute_meta.size(), + /*.mem_buffer =*/ state_.compute_meta.data(), + /*.no_alloc =*/ true, + }; + + struct ggml_context * ctx0 = ggml_init(params); + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false); + + struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_name(inp_tokens, "inp_text_tokens"); + ggml_set_input(inp_tokens); + + struct ggml_tensor * cur = ggml_get_rows(ctx0, model_.text_embd, inp_tokens); + cur = ggml_mul_mat(ctx0, model_.text_proj_fc1, cur); + cur = ggml_add(ctx0, cur, model_.text_proj_fc1_bias); + cur = ggml_silu(ctx0, cur); + cur = ggml_mul_mat(ctx0, model_.text_proj_fc2, cur); + cur = ggml_add(ctx0, cur, model_.text_proj_fc2_bias); + + ggml_set_name(cur, "text_proj_out"); + ggml_set_output(cur); + ggml_build_forward_expand(gf, cur); + + if (!ggml_backend_sched_alloc_graph(state_.sched, gf)) { + error_msg_ = "Failed to allocate text projection graph"; + ggml_free(ctx0); + return false; + } + + struct ggml_tensor * inp = ggml_graph_get_tensor(gf, "inp_text_tokens"); + ggml_backend_tensor_set(inp, text_tokens, 0, n_tokens * sizeof(int32_t)); + + if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) { + error_msg_ = "Failed to compute text projection graph"; + ggml_backend_sched_reset(state_.sched); + ggml_free(ctx0); + return false; + } + + struct ggml_tensor * out = ggml_graph_get_tensor(gf, "text_proj_out"); + if (!out) { + error_msg_ = "Failed to find text projection output tensor"; + ggml_backend_sched_reset(state_.sched); + ggml_free(ctx0); + return false; + } + + output.resize((size_t)model_.config.hidden_size * n_tokens); + ggml_backend_tensor_get(out, output.data(), 0, output.size() * sizeof(float)); + + ggml_backend_sched_reset(state_.sched); + ggml_free(ctx0); + return true; +} + +bool TTSTransformer::build_prefill_graph(const int32_t * text_tokens, int32_t n_tokens, + const float * speaker_embd, int32_t language_id, + std::vector & prefill_embd, + std::vector & trailing_text_hidden, + std::vector & tts_pad_embed) { + if (!text_tokens) { + error_msg_ = "text_tokens is null"; + return false; + } + if (n_tokens < 4) { + error_msg_ = "Need at least 4 text tokens for prefill"; + return false; + } + + const auto & cfg = model_.config; + const int32_t hidden_size = cfg.hidden_size; + + int32_t special_tokens[3] = { + cfg.tts_bos_token_id, + cfg.tts_eos_token_id, + cfg.tts_pad_token_id, + }; + + std::vector special_proj; + if (!project_text_tokens(special_tokens, 3, special_proj)) { + return false; + } + + std::vector tts_bos_embed(hidden_size); + std::vector tts_eos_embed(hidden_size); + tts_pad_embed.resize(hidden_size); + memcpy(tts_bos_embed.data(), special_proj.data() + 0 * hidden_size, hidden_size * sizeof(float)); + memcpy(tts_eos_embed.data(), special_proj.data() + 1 * hidden_size, hidden_size * sizeof(float)); + memcpy(tts_pad_embed.data(), special_proj.data() + 2 * hidden_size, hidden_size * sizeof(float)); + + std::vector role_embed; + if (!project_text_tokens(text_tokens, 3, role_embed)) { + return false; + } + + std::vector codec_prefill_tokens; + if (language_id < 0) { + codec_prefill_tokens = { + cfg.codec_nothink_id, + cfg.codec_think_bos_id, + cfg.codec_think_eos_id, + }; + } else { + codec_prefill_tokens = { + cfg.codec_think_id, + cfg.codec_think_bos_id, + language_id, + cfg.codec_think_eos_id, + }; + } + + std::vector codec_prefill_embed; + if (!lookup_embedding_rows(model_.codec_embd, codec_prefill_tokens.data(), + (int32_t)codec_prefill_tokens.size(), + "inp_codec_prefill_tokens", "codec_prefill_rows", + codec_prefill_embed)) { + return false; + } + + int32_t codec_tail_tokens[2] = { cfg.codec_pad_id, cfg.codec_bos_id }; + std::vector codec_tail_embed; + if (!lookup_embedding_rows(model_.codec_embd, codec_tail_tokens, 2, + "inp_codec_tail_tokens", "codec_tail_rows", + codec_tail_embed)) { + return false; + } + + const bool has_speaker = (speaker_embd != nullptr); + const int32_t codec_input_len = (int32_t)codec_prefill_tokens.size() + (has_speaker ? 1 : 0) + 2; + std::vector codec_input_embedding((size_t)codec_input_len * hidden_size); + + int32_t dst_token = 0; + memcpy(codec_input_embedding.data(), codec_prefill_embed.data(), codec_prefill_embed.size() * sizeof(float)); + dst_token += (int32_t)codec_prefill_tokens.size(); + + if (has_speaker) { + memcpy(codec_input_embedding.data() + (size_t)dst_token * hidden_size, + speaker_embd, hidden_size * sizeof(float)); + ++dst_token; + } + + memcpy(codec_input_embedding.data() + (size_t)dst_token * hidden_size, + codec_tail_embed.data(), codec_tail_embed.size() * sizeof(float)); + + const int32_t codec_plus_overlay_len = codec_input_len - 1; + std::vector codec_plus_overlay((size_t)codec_plus_overlay_len * hidden_size); + for (int32_t t = 0; t < codec_plus_overlay_len; ++t) { + const float * overlay = (t == codec_plus_overlay_len - 1) + ? tts_bos_embed.data() + : tts_pad_embed.data(); + const float * codec_row = codec_input_embedding.data() + (size_t)t * hidden_size; + float * out_row = codec_plus_overlay.data() + (size_t)t * hidden_size; + for (int32_t h = 0; h < hidden_size; ++h) { + out_row[h] = overlay[h] + codec_row[h]; + } + } + + std::vector first_text_embed; + if (!project_text_tokens(text_tokens + 3, 1, first_text_embed)) { + return false; + } + + std::vector first_text_plus_codec_bos(hidden_size); + const float * codec_bos_embed = codec_input_embedding.data() + (size_t)(codec_input_len - 1) * hidden_size; + for (int32_t h = 0; h < hidden_size; ++h) { + first_text_plus_codec_bos[h] = first_text_embed[h] + codec_bos_embed[h]; + } + + const int32_t prefill_len = 3 + codec_plus_overlay_len + 1; + prefill_embd.resize((size_t)prefill_len * hidden_size); + memcpy(prefill_embd.data(), role_embed.data(), role_embed.size() * sizeof(float)); + memcpy(prefill_embd.data() + (size_t)3 * hidden_size, + codec_plus_overlay.data(), codec_plus_overlay.size() * sizeof(float)); + memcpy(prefill_embd.data() + (size_t)(prefill_len - 1) * hidden_size, + first_text_plus_codec_bos.data(), hidden_size * sizeof(float)); + + const int32_t trailing_token_count = std::max(0, n_tokens - 9); + std::vector trailing_text_proj; + if (trailing_token_count > 0) { + if (!project_text_tokens(text_tokens + 4, trailing_token_count, trailing_text_proj)) { + return false; + } + } + + const int32_t trailing_len = trailing_token_count + 1; + trailing_text_hidden.resize((size_t)trailing_len * hidden_size); + if (trailing_token_count > 0) { + memcpy(trailing_text_hidden.data(), trailing_text_proj.data(), trailing_text_proj.size() * sizeof(float)); + } + memcpy(trailing_text_hidden.data() + (size_t)(trailing_len - 1) * hidden_size, + tts_eos_embed.data(), hidden_size * sizeof(float)); + + return true; +} + +struct ggml_cgraph * TTSTransformer::build_prefill_forward_graph(int32_t n_tokens, int32_t n_past) { + const auto & cfg = model_.config; + const int n_head = cfg.n_attention_heads; + const int n_kv_head = cfg.n_key_value_heads; + const int head_dim = cfg.head_dim; + const int hidden_size = cfg.hidden_size; + const float eps = cfg.rms_norm_eps; + const float rope_theta = cfg.rope_theta; + const int n_layer = cfg.n_layers; + + struct ggml_init_params params = { + /*.mem_size =*/ state_.compute_meta.size(), + /*.mem_buffer =*/ state_.compute_meta.data(), + /*.no_alloc =*/ true, + }; + + struct ggml_context * ctx0 = ggml_init(params); + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false); + + struct ggml_tensor * inp_prefill_embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hidden_size, n_tokens); + ggml_set_name(inp_prefill_embd, "inp_prefill_embd"); + ggml_set_input(inp_prefill_embd); + + struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_name(inp_pos, "inp_pos"); + ggml_set_input(inp_pos); + + struct ggml_tensor * cur = inp_prefill_embd; + + struct ggml_tensor * inpL = cur; + + const float KQscale = 1.0f / sqrtf(float(head_dim)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model_.layers[il]; + + cur = ggml_rms_norm(ctx0, inpL, eps); + cur = ggml_mul(ctx0, cur, layer.attn_norm); + + struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur); + struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur); + struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur); + + Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, head_dim, n_kv_head, n_tokens); + + if (layer.attn_q_norm) { + Qcur = ggml_rms_norm(ctx0, Qcur, eps); + Qcur = ggml_mul(ctx0, Qcur, layer.attn_q_norm); + } + + if (layer.attn_k_norm) { + Kcur = ggml_rms_norm(ctx0, Kcur, eps); + Kcur = ggml_mul(ctx0, Kcur, layer.attn_k_norm); + } + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + struct ggml_tensor * k_cache = state_.cache.k_cache[il]; + struct ggml_tensor * v_cache = state_.cache.v_cache[il]; + + struct ggml_tensor * k_cache_view = ggml_view_3d(ctx0, k_cache, + head_dim, n_kv_head, n_tokens, + k_cache->nb[1], k_cache->nb[2], + n_past * k_cache->nb[2]); + + struct ggml_tensor * v_cache_view = ggml_view_3d(ctx0, v_cache, + head_dim, n_kv_head, n_tokens, + v_cache->nb[1], v_cache->nb[2], + n_past * v_cache->nb[2]); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Kcur, k_cache_view)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Vcur, v_cache_view)); + + int n_kv = n_past + n_tokens; + + struct ggml_tensor * K = ggml_view_3d(ctx0, k_cache, + head_dim, n_kv_head, n_kv, + k_cache->nb[1], k_cache->nb[2], 0); + + struct ggml_tensor * V = ggml_view_3d(ctx0, v_cache, + head_dim, n_kv_head, n_kv, + v_cache->nb[1], v_cache->nb[2], 0); + + struct ggml_tensor * Q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + K = ggml_permute(ctx0, K, 0, 2, 1, 3); + V = ggml_permute(ctx0, V, 0, 2, 1, 3); + + struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); + KQ = ggml_scale(ctx0, KQ, KQscale); + KQ = ggml_diag_mask_inf(ctx0, KQ, n_past); + KQ = ggml_soft_max(ctx0, KQ); + + V = ggml_cont(ctx0, ggml_transpose(ctx0, V)); + + struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ); + KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3); + cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, n_tokens); + + cur = ggml_mul_mat(ctx0, layer.attn_output, cur); + cur = ggml_add(ctx0, cur, inpL); + struct ggml_tensor * inpFF = cur; + + cur = ggml_rms_norm(ctx0, inpFF, eps); + cur = ggml_mul(ctx0, cur, layer.ffn_norm); + + struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur); + struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur); + + gate = ggml_silu(ctx0, gate); + + cur = ggml_mul(ctx0, gate, up); + + struct ggml_tensor * ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32); + cur = ggml_mul_mat(ctx0, ffn_down_f32, cur); + + inpL = ggml_add(ctx0, cur, inpFF); + } + + cur = inpL; + + cur = ggml_rms_norm(ctx0, cur, eps); + cur = ggml_mul(ctx0, cur, model_.output_norm); + ggml_set_name(cur, "hidden_states"); + ggml_set_output(cur); + + struct ggml_tensor * logits = ggml_mul_mat(ctx0, model_.codec_head, cur); + ggml_set_name(logits, "logits"); + ggml_set_output(logits); + + ggml_build_forward_expand(gf, logits); + + ggml_free(ctx0); + + return gf; +} + +struct ggml_cgraph * TTSTransformer::build_step_graph(int32_t n_past) { + const auto & cfg = model_.config; + const int n_head = cfg.n_attention_heads; + const int n_kv_head = cfg.n_key_value_heads; + const int head_dim = cfg.head_dim; + const int hidden_size = cfg.hidden_size; + const float eps = cfg.rms_norm_eps; + const float rope_theta = cfg.rope_theta; + const int n_layer = cfg.n_layers; + const int n_tokens = 1; + + struct ggml_init_params params = { + /*.mem_size =*/ state_.compute_meta.size(), + /*.mem_buffer =*/ state_.compute_meta.data(), + /*.no_alloc =*/ true, + }; + + struct ggml_context * ctx0 = ggml_init(params); + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false); + + struct ggml_tensor * inp_step_embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hidden_size, 1); + ggml_set_name(inp_step_embd, "inp_step_embd"); + ggml_set_input(inp_step_embd); + + struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1); + ggml_set_name(inp_pos, "inp_pos"); + ggml_set_input(inp_pos); + + struct ggml_tensor * cur = inp_step_embd; + + struct ggml_tensor * inpL = cur; + + const float KQscale = 1.0f / sqrtf(float(head_dim)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model_.layers[il]; + + cur = ggml_rms_norm(ctx0, inpL, eps); + cur = ggml_mul(ctx0, cur, layer.attn_norm); + + struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur); + struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur); + struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur); + + Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, head_dim, n_kv_head, n_tokens); + + if (layer.attn_q_norm) { + Qcur = ggml_rms_norm(ctx0, Qcur, eps); + Qcur = ggml_mul(ctx0, Qcur, layer.attn_q_norm); + } + + if (layer.attn_k_norm) { + Kcur = ggml_rms_norm(ctx0, Kcur, eps); + Kcur = ggml_mul(ctx0, Kcur, layer.attn_k_norm); + } + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + struct ggml_tensor * k_cache = state_.cache.k_cache[il]; + struct ggml_tensor * v_cache = state_.cache.v_cache[il]; + + struct ggml_tensor * k_cache_view = ggml_view_3d(ctx0, k_cache, + head_dim, n_kv_head, n_tokens, + k_cache->nb[1], k_cache->nb[2], + n_past * k_cache->nb[2]); + + struct ggml_tensor * v_cache_view = ggml_view_3d(ctx0, v_cache, + head_dim, n_kv_head, n_tokens, + v_cache->nb[1], v_cache->nb[2], + n_past * v_cache->nb[2]); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Kcur, k_cache_view)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Vcur, v_cache_view)); + + int n_kv = n_past + n_tokens; + + struct ggml_tensor * K = ggml_view_3d(ctx0, k_cache, + head_dim, n_kv_head, n_kv, + k_cache->nb[1], k_cache->nb[2], 0); + + struct ggml_tensor * V = ggml_view_3d(ctx0, v_cache, + head_dim, n_kv_head, n_kv, + v_cache->nb[1], v_cache->nb[2], 0); + + struct ggml_tensor * Q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + K = ggml_permute(ctx0, K, 0, 2, 1, 3); + V = ggml_permute(ctx0, V, 0, 2, 1, 3); + + struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); + KQ = ggml_scale(ctx0, KQ, KQscale); + KQ = ggml_diag_mask_inf(ctx0, KQ, n_past); + KQ = ggml_soft_max(ctx0, KQ); + + V = ggml_cont(ctx0, ggml_transpose(ctx0, V)); + + struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ); + KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3); + cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, n_tokens); + + cur = ggml_mul_mat(ctx0, layer.attn_output, cur); + cur = ggml_add(ctx0, cur, inpL); + struct ggml_tensor * inpFF = cur; + + cur = ggml_rms_norm(ctx0, inpFF, eps); + cur = ggml_mul(ctx0, cur, layer.ffn_norm); + + struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur); + struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur); + + gate = ggml_silu(ctx0, gate); + + cur = ggml_mul(ctx0, gate, up); + + struct ggml_tensor * ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32); + cur = ggml_mul_mat(ctx0, ffn_down_f32, cur); + + inpL = ggml_add(ctx0, cur, inpFF); + } + + cur = inpL; + + cur = ggml_rms_norm(ctx0, cur, eps); + cur = ggml_mul(ctx0, cur, model_.output_norm); + ggml_set_name(cur, "hidden_states"); + ggml_set_output(cur); + + struct ggml_tensor * logits = ggml_mul_mat(ctx0, model_.codec_head, cur); + ggml_set_name(logits, "logits"); + ggml_set_output(logits); + + ggml_build_forward_expand(gf, logits); + + ggml_free(ctx0); + + return gf; +} + +struct ggml_cgraph * TTSTransformer::build_code_pred_graph(int32_t n_prev_codes) { + const auto & cfg = model_.config; + const int n_head = cfg.n_attention_heads; + const int n_kv_head = cfg.n_key_value_heads; + const int head_dim = cfg.head_dim; + const int hidden_size = cfg.hidden_size; + const float eps = cfg.rms_norm_eps; + const int n_layer = cfg.code_pred_layers; + const int n_codebooks = cfg.n_codebooks; + + struct ggml_init_params params = { + /*.mem_size =*/ state_.compute_meta.size(), + /*.mem_buffer =*/ state_.compute_meta.data(), + /*.no_alloc =*/ true, + }; + + struct ggml_context * ctx0 = ggml_init(params); + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false); + + struct ggml_tensor * inp_hidden = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hidden_size); + ggml_set_name(inp_hidden, "inp_hidden"); + ggml_set_input(inp_hidden); + + struct ggml_tensor * inp_prev_codes = nullptr; + if (n_prev_codes > 0) { + inp_prev_codes = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_prev_codes); + ggml_set_name(inp_prev_codes, "inp_prev_codes"); + ggml_set_input(inp_prev_codes); + } + + struct ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_hidden, hidden_size, 1); + + if (n_prev_codes > 0 && inp_prev_codes) { + for (int cb = 0; cb < n_prev_codes && cb < n_codebooks - 1; ++cb) { + struct ggml_tensor * code_idx = ggml_view_1d(ctx0, inp_prev_codes, 1, cb * sizeof(int32_t)); + struct ggml_tensor * code_embd = ggml_get_rows(ctx0, model_.code_pred_embd[cb], code_idx); + cur = ggml_add(ctx0, cur, code_embd); + } + } + + struct ggml_tensor * inpL = cur; + + const float KQscale = 1.0f / sqrtf(float(head_dim)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model_.code_pred_layers[il]; + + cur = ggml_rms_norm(ctx0, inpL, eps); + cur = ggml_mul(ctx0, cur, layer.attn_norm); + + struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur); + struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur); + struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur); + + Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, 1); + Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, 1); + Vcur = ggml_reshape_3d(ctx0, Vcur, head_dim, n_kv_head, 1); + + if (layer.attn_q_norm) { + Qcur = ggml_rms_norm(ctx0, Qcur, eps); + Qcur = ggml_mul(ctx0, Qcur, layer.attn_q_norm); + } + + if (layer.attn_k_norm) { + Kcur = ggml_rms_norm(ctx0, Kcur, eps); + Kcur = ggml_mul(ctx0, Kcur, layer.attn_k_norm); + } + + struct ggml_tensor * Q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + struct ggml_tensor * K = ggml_permute(ctx0, Kcur, 0, 2, 1, 3); + struct ggml_tensor * V = ggml_permute(ctx0, Vcur, 0, 2, 1, 3); + + struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); + KQ = ggml_scale(ctx0, KQ, KQscale); + KQ = ggml_soft_max(ctx0, KQ); + + V = ggml_cont(ctx0, ggml_transpose(ctx0, V)); + + struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ); + KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3); + cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, 1); + + cur = ggml_mul_mat(ctx0, layer.attn_output, cur); + cur = ggml_add(ctx0, cur, inpL); + struct ggml_tensor * inpFF = cur; + + cur = ggml_rms_norm(ctx0, inpFF, eps); + cur = ggml_mul(ctx0, cur, layer.ffn_norm); + + struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur); + struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur); + + gate = ggml_silu(ctx0, gate); + + cur = ggml_mul(ctx0, gate, up); + + struct ggml_tensor * old_ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32); + cur = ggml_mul_mat(ctx0, old_ffn_down_f32, cur); + + inpL = ggml_add(ctx0, cur, inpFF); + } + + cur = inpL; + + std::vector all_logits; + for (int cb = 0; cb < n_codebooks - 1; ++cb) { + struct ggml_tensor * cb_logits = ggml_mul_mat(ctx0, model_.code_pred_head[cb], cur); + ggml_format_name(cb_logits, "logits_cb%d", cb + 1); + ggml_set_output(cb_logits); + all_logits.push_back(cb_logits); + } + + for (auto * logits : all_logits) { + ggml_build_forward_expand(gf, logits); + } + + ggml_free(ctx0); + + return gf; +} + +struct ggml_cgraph * TTSTransformer::build_code_pred_prefill_graph() { + const auto & cfg = model_.config; + const int n_head = cfg.n_attention_heads; + const int n_kv_head = cfg.n_key_value_heads; + const int head_dim = cfg.head_dim; + const int hidden_size = cfg.hidden_size; + const float eps = cfg.rms_norm_eps; + const float rope_theta = cfg.rope_theta; + const int n_layer = cfg.code_pred_layers; + const int n_tokens = 2; + + struct ggml_init_params params = { + /*.mem_size =*/ state_.compute_meta.size(), + /*.mem_buffer =*/ state_.compute_meta.data(), + /*.no_alloc =*/ true, + }; + + struct ggml_context * ctx0 = ggml_init(params); + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false); + + // Input: past_hidden from talker [hidden_size] + struct ggml_tensor * inp_hidden = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hidden_size); + ggml_set_name(inp_hidden, "inp_hidden"); + ggml_set_input(inp_hidden); + + // Input: codebook 0 token embedding [hidden_size] (pre-computed using talker's codec_embd) + struct ggml_tensor * inp_cb0_embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hidden_size); + ggml_set_name(inp_cb0_embd, "inp_cb0_embd"); + ggml_set_input(inp_cb0_embd); + + struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_name(inp_pos, "inp_pos"); + ggml_set_input(inp_pos); + + // Concatenate [past_hidden, cb0_embd] -> [2, hidden_size] + struct ggml_tensor * hidden_2d = ggml_reshape_2d(ctx0, inp_hidden, hidden_size, 1); + struct ggml_tensor * cb0_2d = ggml_reshape_2d(ctx0, inp_cb0_embd, hidden_size, 1); + struct ggml_tensor * cur = ggml_concat(ctx0, hidden_2d, cb0_2d, 1); + + struct ggml_tensor * inpL = cur; + + const float KQscale = 1.0f / sqrtf(float(head_dim)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model_.code_pred_layers[il]; + + cur = ggml_rms_norm(ctx0, inpL, eps); + cur = ggml_mul(ctx0, cur, layer.attn_norm); + + struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur); + struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur); + struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur); + + Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, head_dim, n_kv_head, n_tokens); + + if (layer.attn_q_norm) { + Qcur = ggml_rms_norm(ctx0, Qcur, eps); + Qcur = ggml_mul(ctx0, Qcur, layer.attn_q_norm); + } + + if (layer.attn_k_norm) { + Kcur = ggml_rms_norm(ctx0, Kcur, eps); + Kcur = ggml_mul(ctx0, Kcur, layer.attn_k_norm); + } + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + struct ggml_tensor * k_cache = state_.code_pred_cache.k_cache[il]; + struct ggml_tensor * v_cache = state_.code_pred_cache.v_cache[il]; + + // Store at position 0 (prefill starts fresh) + struct ggml_tensor * k_cache_view = ggml_view_3d(ctx0, k_cache, + head_dim, n_kv_head, n_tokens, + k_cache->nb[1], k_cache->nb[2], 0); + + struct ggml_tensor * v_cache_view = ggml_view_3d(ctx0, v_cache, + head_dim, n_kv_head, n_tokens, + v_cache->nb[1], v_cache->nb[2], 0); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Kcur, k_cache_view)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Vcur, v_cache_view)); + + struct ggml_tensor * Q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + struct ggml_tensor * K = ggml_permute(ctx0, Kcur, 0, 2, 1, 3); + struct ggml_tensor * V = ggml_permute(ctx0, Vcur, 0, 2, 1, 3); + + struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); + KQ = ggml_scale(ctx0, KQ, KQscale); + KQ = ggml_diag_mask_inf(ctx0, KQ, 0); + KQ = ggml_soft_max(ctx0, KQ); + + V = ggml_cont(ctx0, ggml_transpose(ctx0, V)); + + struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ); + KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3); + cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, n_tokens); + + cur = ggml_mul_mat(ctx0, layer.attn_output, cur); + cur = ggml_add(ctx0, cur, inpL); + struct ggml_tensor * inpFF = cur; + + cur = ggml_rms_norm(ctx0, inpFF, eps); + cur = ggml_mul(ctx0, cur, layer.ffn_norm); + + struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur); + struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur); + + gate = ggml_silu(ctx0, gate); + + cur = ggml_mul(ctx0, gate, up); + + struct ggml_tensor * ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32); + cur = ggml_mul_mat(ctx0, ffn_down_f32, cur); + + inpL = ggml_add(ctx0, cur, inpFF); + } + + cur = inpL; + + cur = ggml_rms_norm(ctx0, cur, eps); + cur = ggml_mul(ctx0, cur, model_.code_pred_output_norm); + + struct ggml_tensor * last_hidden = ggml_view_2d(ctx0, cur, hidden_size, 1, + cur->nb[1], hidden_size * sizeof(float)); + + struct ggml_tensor * logits = ggml_mul_mat(ctx0, model_.code_pred_head[0], last_hidden); + ggml_set_name(logits, "logits"); + ggml_set_output(logits); + + ggml_build_forward_expand(gf, logits); + + ggml_free(ctx0); + + return gf; +} + +struct ggml_cgraph * TTSTransformer::build_code_pred_step_graph(int32_t n_past, int32_t generation_step) { + const auto & cfg = model_.config; + const int n_head = cfg.n_attention_heads; + const int n_kv_head = cfg.n_key_value_heads; + const int head_dim = cfg.head_dim; + const int hidden_size = cfg.hidden_size; + const float eps = cfg.rms_norm_eps; + const float rope_theta = cfg.rope_theta; + const int n_layer = cfg.code_pred_layers; + const int n_tokens = 1; + + struct ggml_init_params params = { + /*.mem_size =*/ state_.compute_meta.size(), + /*.mem_buffer =*/ state_.compute_meta.data(), + /*.no_alloc =*/ true, + }; + + struct ggml_context * ctx0 = ggml_init(params); + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false); + + struct ggml_tensor * inp_hidden = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hidden_size); + ggml_set_name(inp_hidden, "inp_hidden"); + ggml_set_input(inp_hidden); + + struct ggml_tensor * inp_code = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1); + ggml_set_name(inp_code, "inp_code"); + ggml_set_input(inp_code); + + struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1); + ggml_set_name(inp_pos, "inp_pos"); + ggml_set_input(inp_pos); + + struct ggml_tensor * cur; + if (generation_step == 0) { + cur = ggml_reshape_2d(ctx0, inp_hidden, hidden_size, 1); + } else { + cur = ggml_get_rows(ctx0, model_.code_pred_embd[generation_step - 1], inp_code); + cur = ggml_reshape_2d(ctx0, cur, hidden_size, 1); + } + + struct ggml_tensor * inpL = cur; + + const float KQscale = 1.0f / sqrtf(float(head_dim)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model_.code_pred_layers[il]; + + cur = ggml_rms_norm(ctx0, inpL, eps); + cur = ggml_mul(ctx0, cur, layer.attn_norm); + + struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur); + struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur); + struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur); + + Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, head_dim, n_kv_head, n_tokens); + + if (layer.attn_q_norm) { + Qcur = ggml_rms_norm(ctx0, Qcur, eps); + Qcur = ggml_mul(ctx0, Qcur, layer.attn_q_norm); + } + + if (layer.attn_k_norm) { + Kcur = ggml_rms_norm(ctx0, Kcur, eps); + Kcur = ggml_mul(ctx0, Kcur, layer.attn_k_norm); + } + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, + head_dim, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + struct ggml_tensor * k_cache = state_.code_pred_cache.k_cache[il]; + struct ggml_tensor * v_cache = state_.code_pred_cache.v_cache[il]; + + struct ggml_tensor * k_cache_view = ggml_view_3d(ctx0, k_cache, + head_dim, n_kv_head, n_tokens, + k_cache->nb[1], k_cache->nb[2], + n_past * k_cache->nb[2]); + + struct ggml_tensor * v_cache_view = ggml_view_3d(ctx0, v_cache, + head_dim, n_kv_head, n_tokens, + v_cache->nb[1], v_cache->nb[2], + n_past * v_cache->nb[2]); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Kcur, k_cache_view)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Vcur, v_cache_view)); + + int n_kv = n_past + n_tokens; + + struct ggml_tensor * K = ggml_view_3d(ctx0, k_cache, + head_dim, n_kv_head, n_kv, + k_cache->nb[1], k_cache->nb[2], 0); + + struct ggml_tensor * V = ggml_view_3d(ctx0, v_cache, + head_dim, n_kv_head, n_kv, + v_cache->nb[1], v_cache->nb[2], 0); + + struct ggml_tensor * Q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + K = ggml_permute(ctx0, K, 0, 2, 1, 3); + V = ggml_permute(ctx0, V, 0, 2, 1, 3); + + struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); + KQ = ggml_scale(ctx0, KQ, KQscale); + KQ = ggml_diag_mask_inf(ctx0, KQ, n_past); + KQ = ggml_soft_max(ctx0, KQ); + + V = ggml_cont(ctx0, ggml_transpose(ctx0, V)); + + struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ); + KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3); + cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, n_tokens); + + cur = ggml_mul_mat(ctx0, layer.attn_output, cur); + cur = ggml_add(ctx0, cur, inpL); + struct ggml_tensor * inpFF = cur; + + cur = ggml_rms_norm(ctx0, inpFF, eps); + cur = ggml_mul(ctx0, cur, layer.ffn_norm); + + struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur); + struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur); + + gate = ggml_silu(ctx0, gate); + + cur = ggml_mul(ctx0, gate, up); + + struct ggml_tensor * step_ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32); + cur = ggml_mul_mat(ctx0, step_ffn_down_f32, cur); + + inpL = ggml_add(ctx0, cur, inpFF); + } + + cur = inpL; + + cur = ggml_rms_norm(ctx0, cur, eps); + cur = ggml_mul(ctx0, cur, model_.code_pred_output_norm); + + struct ggml_tensor * logits = ggml_mul_mat(ctx0, model_.code_pred_head[generation_step], cur); + ggml_set_name(logits, "logits"); + ggml_set_output(logits); + + ggml_build_forward_expand(gf, logits); + + ggml_free(ctx0); + + return gf; +} + +bool TTSTransformer::forward_prefill(const float * prefill_embd, int32_t n_tokens, + int32_t n_past, std::vector & output, + std::vector * logits_out) { + if (!model_.ctx) { + error_msg_ = "Model not loaded"; + return false; + } + if (!prefill_embd) { + error_msg_ = "prefill_embd is null"; + return false; + } + if (n_tokens <= 0) { + error_msg_ = "n_tokens must be > 0"; + return false; + } + + if (state_.cache.n_ctx == 0) { + const int32_t min_ctx = std::max(256, n_past + n_tokens + 16); + if (!init_kv_cache(min_ctx)) { + return false; + } + } + + if (n_past + n_tokens > state_.cache.n_ctx) { + error_msg_ = "Context length exceeded"; + return false; + } + +#ifdef QWEN3_TTS_TIMING + using clk = std::chrono::high_resolution_clock; + auto t0 = clk::now(), t1 = t0; +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + struct ggml_cgraph * gf = build_prefill_forward_graph(n_tokens, n_past); +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_prefill_graph_build_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (!ggml_backend_sched_alloc_graph(state_.sched, gf)) { + error_msg_ = "Failed to allocate graph"; + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_prefill_graph_alloc_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + struct ggml_tensor * inp_prefill = ggml_graph_get_tensor(gf, "inp_prefill_embd"); + if (inp_prefill) { + ggml_backend_tensor_set(inp_prefill, prefill_embd, 0, + (size_t)n_tokens * model_.config.hidden_size * sizeof(float)); + } + + struct ggml_tensor * inp_pos = ggml_graph_get_tensor(gf, "inp_pos"); + if (inp_pos) { + std::vector positions(n_tokens); + for (int i = 0; i < n_tokens; ++i) { + positions[i] = n_past + i; + } + ggml_backend_tensor_set(inp_pos, positions.data(), 0, n_tokens * sizeof(int32_t)); + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_prefill_data_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) { + error_msg_ = "Failed to compute graph"; + ggml_backend_sched_reset(state_.sched); + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_prefill_compute_ms += std::chrono::duration(t1 - t0).count(); +#endif + + struct ggml_tensor * hidden = ggml_graph_get_tensor(gf, "hidden_states"); + if (!hidden) { + error_msg_ = "Failed to find hidden_states tensor"; + ggml_backend_sched_reset(state_.sched); + return false; + } + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + output.resize(n_tokens * model_.config.hidden_size); + ggml_backend_tensor_get(hidden, output.data(), 0, output.size() * sizeof(float)); + + last_hidden_.resize(model_.config.hidden_size); + ggml_backend_tensor_get(hidden, last_hidden_.data(), + (n_tokens - 1) * model_.config.hidden_size * sizeof(float), + model_.config.hidden_size * sizeof(float)); + + if (logits_out) { + struct ggml_tensor * logits = ggml_graph_get_tensor(gf, "logits"); + if (!logits) { + error_msg_ = "Failed to find logits tensor"; + ggml_backend_sched_reset(state_.sched); + return false; + } + + logits_out->resize(model_.config.codec_vocab_size); + ggml_backend_tensor_get(logits, logits_out->data(), + (n_tokens - 1) * model_.config.codec_vocab_size * sizeof(float), + model_.config.codec_vocab_size * sizeof(float)); + } + + state_.cache.n_used = n_past + n_tokens; + + ggml_backend_sched_reset(state_.sched); +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_prefill_data_ms += std::chrono::duration(t1 - t0).count(); +#endif + + return true; +} + +bool TTSTransformer::forward_text(const int32_t * text_tokens, int32_t n_tokens, + const float * speaker_embd, int32_t n_past, + std::vector & output) { + if (!text_tokens) { + error_msg_ = "text_tokens is null"; + return false; + } + if (n_tokens <= 0) { + error_msg_ = "n_tokens must be > 0"; + return false; + } + + std::vector projected; + if (!project_text_tokens(text_tokens, n_tokens, projected)) { + return false; + } + + if (speaker_embd) { + const int32_t hidden_size = model_.config.hidden_size; + for (int32_t t = 0; t < n_tokens; ++t) { + float * row = projected.data() + (size_t)t * hidden_size; + for (int32_t h = 0; h < hidden_size; ++h) { + row[h] += speaker_embd[h]; + } + } + } + + return forward_prefill(projected.data(), n_tokens, n_past, output, nullptr); +} + +bool TTSTransformer::forward_step(const float * step_embd, int32_t n_past, + std::vector & output, + std::vector * hidden_out) { + if (!model_.ctx) { + error_msg_ = "Model not loaded"; + return false; + } + if (!step_embd) { + error_msg_ = "step_embd is null"; + return false; + } + + if (state_.cache.n_ctx == 0) { + const int32_t min_ctx = std::max(256, n_past + 1 + 16); + if (!init_kv_cache(min_ctx)) { + return false; + } + } + + if (n_past + 1 > state_.cache.n_ctx) { + error_msg_ = "Context length exceeded"; + return false; + } + +#ifdef QWEN3_TTS_TIMING + using clk = std::chrono::high_resolution_clock; + auto t0 = clk::now(), t1 = t0; +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + struct ggml_cgraph * gf = build_step_graph(n_past); +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_talker_graph_build_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (!ggml_backend_sched_alloc_graph(state_.sched, gf)) { + error_msg_ = "Failed to allocate graph"; + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_talker_graph_alloc_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + struct ggml_tensor * inp_step = ggml_graph_get_tensor(gf, "inp_step_embd"); + if (inp_step) { + ggml_backend_tensor_set(inp_step, step_embd, 0, + model_.config.hidden_size * sizeof(float)); + } + + struct ggml_tensor * inp_pos = ggml_graph_get_tensor(gf, "inp_pos"); + if (inp_pos) { + int32_t pos = n_past; + ggml_backend_tensor_set(inp_pos, &pos, 0, sizeof(int32_t)); + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_talker_data_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) { + error_msg_ = "Failed to compute graph"; + ggml_backend_sched_reset(state_.sched); + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_talker_compute_ms += std::chrono::duration(t1 - t0).count(); +#endif + + struct ggml_tensor * hidden = ggml_graph_get_tensor(gf, "hidden_states"); + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (hidden) { + last_hidden_.resize(model_.config.hidden_size); + ggml_backend_tensor_get(hidden, last_hidden_.data(), 0, + model_.config.hidden_size * sizeof(float)); + if (hidden_out) { + *hidden_out = last_hidden_; + } + } + + struct ggml_tensor * logits = ggml_graph_get_tensor(gf, "logits"); + if (!logits) { + error_msg_ = "Failed to find logits tensor"; + ggml_backend_sched_reset(state_.sched); + return false; + } + + output.resize(model_.config.codec_vocab_size); + ggml_backend_tensor_get(logits, output.data(), 0, output.size() * sizeof(float)); + + state_.cache.n_used = n_past + 1; + + ggml_backend_sched_reset(state_.sched); +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_talker_data_ms += std::chrono::duration(t1 - t0).count(); +#endif + + return true; +} + +bool TTSTransformer::forward_codec(int32_t codec_token, int32_t n_past, + std::vector & output) { + std::vector codec_row; + if (!lookup_embedding_rows(model_.codec_embd, &codec_token, 1, + "inp_legacy_codec_token", "legacy_codec_row", + codec_row)) { + return false; + } + + return forward_step(codec_row.data(), n_past, output, nullptr); +} + +bool TTSTransformer::get_hidden_states(std::vector & hidden) const { + if (last_hidden_.empty()) { + return false; + } + hidden = last_hidden_; + return true; +} + +bool TTSTransformer::predict_codes(const float * hidden, const int32_t * prev_codes, + std::vector & output) { + if (!model_.ctx) { + error_msg_ = "Model not loaded"; + return false; + } + + const auto & cfg = model_.config; + int n_prev = (prev_codes != nullptr) ? cfg.n_codebooks - 1 : 0; + + struct ggml_cgraph * gf = build_code_pred_graph(n_prev); + + if (!ggml_backend_sched_alloc_graph(state_.sched, gf)) { + error_msg_ = "Failed to allocate code predictor graph"; + return false; + } + + struct ggml_tensor * inp_hidden = ggml_graph_get_tensor(gf, "inp_hidden"); + if (inp_hidden) { + ggml_backend_tensor_set(inp_hidden, hidden, 0, cfg.hidden_size * sizeof(float)); + } + + if (n_prev > 0) { + struct ggml_tensor * inp_prev = ggml_graph_get_tensor(gf, "inp_prev_codes"); + if (inp_prev) { + ggml_backend_tensor_set(inp_prev, prev_codes, 0, n_prev * sizeof(int32_t)); + } + } + + if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) { + error_msg_ = "Failed to compute code predictor graph"; + ggml_backend_sched_reset(state_.sched); + return false; + } + + output.resize((cfg.n_codebooks - 1) * cfg.code_pred_vocab_size); + + for (int cb = 0; cb < cfg.n_codebooks - 1; ++cb) { + char name[32]; + snprintf(name, sizeof(name), "logits_cb%d", cb + 1); + struct ggml_tensor * cb_logits = ggml_graph_get_tensor(gf, name); + if (cb_logits) { + ggml_backend_tensor_get(cb_logits, output.data() + cb * cfg.code_pred_vocab_size, + 0, cfg.code_pred_vocab_size * sizeof(float)); + } + } + + ggml_backend_sched_reset(state_.sched); + + return true; +} + +static int32_t argmax(const float * data, int32_t n) { + int32_t max_idx = 0; + float max_val = data[0]; + for (int32_t i = 1; i < n; ++i) { + if (data[i] > max_val) { + max_val = data[i]; + max_idx = i; + } + } + return max_idx; +} + +bool TTSTransformer::predict_codes_autoregressive_coreml(const float * hidden, + int32_t codebook_0_token, + std::vector & output, + float temperature, + int32_t top_k) { + if (!use_coreml_code_predictor_ || !coreml_code_predictor_.is_loaded()) { + error_msg_ = "CoreML code predictor is not loaded"; + return false; + } + + const auto & cfg = model_.config; + const int32_t n_steps = cfg.n_codebooks - 1; + + output.resize(n_steps); + std::vector logits_data(cfg.code_pred_vocab_size); + std::vector code_probs(cfg.code_pred_vocab_size); + std::vector seq_embd((size_t)16 * cfg.hidden_size, 0.0f); + +#ifdef QWEN3_TTS_TIMING + using clk = std::chrono::high_resolution_clock; + auto t0 = clk::now(), t1 = t0; +#endif + + auto sample_or_argmax = [&](float * logits_ptr, int32_t vocab_size) -> int32_t { + if (temperature <= 0.0f) { + return argmax(logits_ptr, vocab_size); + } + + for (int32_t i = 0; i < vocab_size; ++i) { + logits_ptr[i] /= temperature; + } + + if (top_k > 0 && top_k < vocab_size) { + std::vector> scored(vocab_size); + for (int32_t i = 0; i < vocab_size; ++i) { + scored[i] = {logits_ptr[i], i}; + } + std::partial_sort(scored.begin(), scored.begin() + top_k, scored.end(), + [](const std::pair & a, const std::pair & b) { + return a.first > b.first; + }); + float threshold = scored[top_k - 1].first; + for (int32_t i = 0; i < vocab_size; ++i) { + if (logits_ptr[i] < threshold) { + logits_ptr[i] = -INFINITY; + } + } + } + + float max_logit = *std::max_element(logits_ptr, logits_ptr + vocab_size); + double sum = 0.0; + for (int32_t i = 0; i < vocab_size; ++i) { + code_probs[i] = expf(logits_ptr[i] - max_logit); + sum += code_probs[i]; + } + for (int32_t i = 0; i < vocab_size; ++i) { + code_probs[i] = (float)(code_probs[i] / sum); + } + + std::discrete_distribution dist(code_probs.begin(), code_probs.begin() + vocab_size); + return dist(rng_); + }; + + memcpy(seq_embd.data(), hidden, (size_t)cfg.hidden_size * sizeof(float)); + if (!lookup_single_embedding_row(model_.codec_embd, codebook_0_token, + seq_embd.data() + cfg.hidden_size)) { + return false; + } + +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_init_ms += std::chrono::duration(t1 - t0).count(); +#endif + + for (int32_t step = 0; step < n_steps; ++step) { + if (step > 0) { + float * dst = seq_embd.data() + (size_t)(step + 1) * cfg.hidden_size; + if (!lookup_single_embedding_row(model_.code_pred_embd[step - 1], output[step - 1], dst)) { + return false; + } + } + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (!coreml_code_predictor_.predict_step(step, seq_embd.data(), step + 2, cfg.hidden_size, logits_data)) { + error_msg_ = "CoreML predictor step failed: " + coreml_code_predictor_.get_error(); + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + const double dt_ms = std::chrono::duration(t1 - t0).count(); + if (timing_) timing_->t_code_pred_compute_ms += dt_ms; + if (timing_) timing_->t_code_pred_coreml_ms += dt_ms; +#endif + + if ((int32_t)logits_data.size() != cfg.code_pred_vocab_size) { + error_msg_ = "CoreML predictor returned unexpected logits size"; + return false; + } + output[step] = sample_or_argmax(logits_data.data(), cfg.code_pred_vocab_size); + +#ifdef QWEN3_TTS_TIMING + if (timing_) { + if (step == 0) { + timing_->t_code_pred_prefill_ms += dt_ms; + } else { + timing_->t_code_pred_steps_ms += dt_ms; + } + } +#endif + } + + return true; +} + +bool TTSTransformer::predict_codes_autoregressive(const float * hidden, int32_t codebook_0_token, + std::vector & output, + float temperature, int32_t top_k) { + if (!model_.ctx) { + error_msg_ = "Model not loaded"; + return false; + } + + const auto & cfg = model_.config; + +#ifdef QWEN3_TTS_TIMING + using clk = std::chrono::high_resolution_clock; + auto t0 = clk::now(), t1 = t0; +#endif + + if (use_coreml_code_predictor_ && coreml_code_predictor_.is_loaded()) { + if (predict_codes_autoregressive_coreml(hidden, codebook_0_token, output, temperature, top_k)) { + return true; + } + if (skip_ggml_code_pred_layers_) { + return false; + } + fprintf(stderr, " CoreML code predictor failed, falling back to GGML: %s\n", error_msg_.c_str()); + use_coreml_code_predictor_ = false; + } + + if (state_.code_pred_cache.n_ctx < 16) { + if (!init_code_pred_kv_cache(16)) { + return false; + } + } + clear_code_pred_kv_cache(); + + output.resize(15); + std::vector logits_data(cfg.code_pred_vocab_size); + + std::vector code_probs(cfg.code_pred_vocab_size); + + // Helper lambda: temperature + top-k sampling (or greedy if temperature <= 0) + auto sample_or_argmax = [&](float * logits_ptr, int32_t vocab_size) -> int32_t { + if (temperature <= 0.0f) { + return argmax(logits_ptr, vocab_size); + } + // Temperature scaling + for (int32_t i = 0; i < vocab_size; ++i) { + logits_ptr[i] /= temperature; + } + // Top-k filtering + if (top_k > 0 && top_k < vocab_size) { + std::vector> scored(vocab_size); + for (int32_t i = 0; i < vocab_size; ++i) { + scored[i] = {logits_ptr[i], i}; + } + std::partial_sort(scored.begin(), scored.begin() + top_k, scored.end(), + [](const std::pair & a, const std::pair & b) { + return a.first > b.first; + }); + float threshold = scored[top_k - 1].first; + for (int32_t i = 0; i < vocab_size; ++i) { + if (logits_ptr[i] < threshold) { + logits_ptr[i] = -INFINITY; + } + } + } + // Softmax + float max_logit = *std::max_element(logits_ptr, logits_ptr + vocab_size); + double sum = 0.0; + for (int32_t i = 0; i < vocab_size; ++i) { + code_probs[i] = expf(logits_ptr[i] - max_logit); + sum += code_probs[i]; + } + for (int32_t i = 0; i < vocab_size; ++i) { + code_probs[i] = (float)(code_probs[i] / sum); + } + // Sample + std::discrete_distribution dist(code_probs.begin(), code_probs.begin() + vocab_size); + return dist(rng_); + }; + + std::vector cb0_embd(cfg.hidden_size); + if (!lookup_single_embedding_row(model_.codec_embd, codebook_0_token, cb0_embd.data())) { + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_init_ms += std::chrono::duration(t1 - t0).count(); +#endif + + // Prefill with 2 tokens [past_hidden, cb0_embd] + { +#ifdef QWEN3_TTS_TIMING + auto t_pf_start = clk::now(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + struct ggml_cgraph * gf = build_code_pred_prefill_graph(); +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_graph_build_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (!ggml_backend_sched_alloc_graph(state_.sched, gf)) { + error_msg_ = "Failed to allocate code predictor prefill graph"; + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_graph_alloc_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + struct ggml_tensor * inp_hidden = ggml_graph_get_tensor(gf, "inp_hidden"); + if (inp_hidden) { + ggml_backend_tensor_set(inp_hidden, hidden, 0, cfg.hidden_size * sizeof(float)); + } + + struct ggml_tensor * inp_cb0_embd = ggml_graph_get_tensor(gf, "inp_cb0_embd"); + if (inp_cb0_embd) { + ggml_backend_tensor_set(inp_cb0_embd, cb0_embd.data(), 0, cfg.hidden_size * sizeof(float)); + } + + struct ggml_tensor * inp_pos = ggml_graph_get_tensor(gf, "inp_pos"); + if (inp_pos) { + int32_t positions[2] = {0, 1}; + ggml_backend_tensor_set(inp_pos, positions, 0, 2 * sizeof(int32_t)); + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_data_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) { + error_msg_ = "Failed to compute code predictor prefill graph"; + ggml_backend_sched_reset(state_.sched); + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_compute_ms += std::chrono::duration(t1 - t0).count(); +#endif + + struct ggml_tensor * logits = ggml_graph_get_tensor(gf, "logits"); + if (!logits) { + error_msg_ = "Failed to find logits tensor in prefill"; + ggml_backend_sched_reset(state_.sched); + return false; + } + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + ggml_backend_tensor_get(logits, logits_data.data(), 0, + cfg.code_pred_vocab_size * sizeof(float)); + + output[0] = sample_or_argmax(logits_data.data(), cfg.code_pred_vocab_size); + + ggml_backend_sched_reset(state_.sched); +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_data_ms += std::chrono::duration(t1 - t0).count(); + if (timing_) timing_->t_code_pred_prefill_ms += std::chrono::duration(t1 - t_pf_start).count(); +#endif + } + + // Generate 14 more tokens autoregressively +#ifdef QWEN3_TTS_TIMING + auto t_steps_start = clk::now(); +#endif + for (int step = 1; step < 15; ++step) { + int32_t n_past = step + 1; + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + struct ggml_cgraph * gf = build_code_pred_step_graph(n_past, step); +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_graph_build_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (!ggml_backend_sched_alloc_graph(state_.sched, gf)) { + error_msg_ = "Failed to allocate code predictor step graph"; + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_graph_alloc_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + struct ggml_tensor * inp_hidden = ggml_graph_get_tensor(gf, "inp_hidden"); + if (inp_hidden) { + ggml_backend_tensor_set(inp_hidden, hidden, 0, cfg.hidden_size * sizeof(float)); + } + + struct ggml_tensor * inp_code = ggml_graph_get_tensor(gf, "inp_code"); + if (inp_code) { + int32_t prev_code = output[step - 1]; + ggml_backend_tensor_set(inp_code, &prev_code, 0, sizeof(int32_t)); + } + + struct ggml_tensor * inp_pos = ggml_graph_get_tensor(gf, "inp_pos"); + if (inp_pos) { + int32_t pos = n_past; + ggml_backend_tensor_set(inp_pos, &pos, 0, sizeof(int32_t)); + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_data_ms += std::chrono::duration(t1 - t0).count(); +#endif + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) { + error_msg_ = "Failed to compute code predictor step graph"; + ggml_backend_sched_reset(state_.sched); + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_compute_ms += std::chrono::duration(t1 - t0).count(); +#endif + + struct ggml_tensor * logits = ggml_graph_get_tensor(gf, "logits"); + if (!logits) { + error_msg_ = "Failed to find logits tensor"; + ggml_backend_sched_reset(state_.sched); + return false; + } + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + ggml_backend_tensor_get(logits, logits_data.data(), 0, + cfg.code_pred_vocab_size * sizeof(float)); + + output[step] = sample_or_argmax(logits_data.data(), cfg.code_pred_vocab_size); + + ggml_backend_sched_reset(state_.sched); +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + if (timing_) timing_->t_code_pred_data_ms += std::chrono::duration(t1 - t0).count(); +#endif + } +#ifdef QWEN3_TTS_TIMING + if (timing_) timing_->t_code_pred_steps_ms += std::chrono::duration(clk::now() - t_steps_start).count(); +#endif + + return true; +} + +bool TTSTransformer::generate(const int32_t * text_tokens, int32_t n_tokens, + const float * speaker_embd, int32_t max_len, + std::vector & output, + int32_t language_id, + float repetition_penalty, + float temperature, + int32_t top_k) { +#ifdef QWEN3_TTS_TIMING + using clk = std::chrono::high_resolution_clock; + tts_timing timing = {}; + auto t_gen_start = clk::now(); + auto t0 = t_gen_start, t1 = t_gen_start; + timing_ = &timing; +#endif + + if (!model_.ctx) { + error_msg_ = "Model not loaded"; + return false; + } + if (!text_tokens) { + error_msg_ = "text_tokens is null"; + return false; + } + if (n_tokens < 4) { + error_msg_ = "Need at least 4 text tokens for generation"; + return false; + } + if (max_len <= 0) { + output.clear(); + return true; + } + + const auto & cfg = model_.config; + + std::vector prefill_embd; + std::vector trailing_text_hidden; + std::vector tts_pad_embed; + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (!build_prefill_graph(text_tokens, n_tokens, speaker_embd, language_id, + prefill_embd, trailing_text_hidden, tts_pad_embed)) { + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + timing.t_prefill_build_ms = std::chrono::duration(t1 - t0).count(); +#endif + + const int32_t prefill_len = (int32_t)(prefill_embd.size() / cfg.hidden_size); + const int32_t trailing_len = (int32_t)(trailing_text_hidden.size() / cfg.hidden_size); + + const int32_t required_ctx = prefill_len + max_len + 8; + if (state_.cache.n_ctx < required_ctx || state_.cache.n_ctx > std::max(required_ctx * 2, 512)) { + if (!init_kv_cache(required_ctx)) { + return false; + } + } + clear_kv_cache(); + + std::vector hidden_out; + std::vector logits; + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (!forward_prefill(prefill_embd.data(), prefill_len, 0, hidden_out, &logits)) { + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + timing.t_prefill_forward_ms = std::chrono::duration(t1 - t0).count(); +#endif + + output.clear(); + output.reserve(max_len * cfg.n_codebooks); + + int32_t n_past = prefill_len; + std::vector frame_codes(cfg.n_codebooks); + std::unordered_set generated_cb0_tokens; + const int32_t suppress_start = cfg.codec_vocab_size - 1024; + + std::vector probs(cfg.codec_vocab_size); + std::vector step_embd(cfg.hidden_size, 0.0f); + std::vector embd_row(cfg.hidden_size); + + for (int frame = 0; frame < max_len; ++frame) { + // Suppress tokens in [codec_vocab_size - 1024, codec_vocab_size), except codec_eos_id + for (int32_t i = suppress_start; i < cfg.codec_vocab_size; ++i) { + if (i != cfg.codec_eos_id) { + logits[i] = -INFINITY; + } + } + + // Repetition penalty (HuggingFace style) on previously generated CB0 tokens + if (repetition_penalty != 1.0f) { + for (int32_t tok : generated_cb0_tokens) { + if (tok >= 0 && tok < cfg.codec_vocab_size) { + if (logits[tok] > 0.0f) { + logits[tok] /= repetition_penalty; + } else { + logits[tok] *= repetition_penalty; + } + } + } + } + + int32_t next_token; + if (temperature <= 0.0f) { + next_token = argmax(logits.data(), cfg.codec_vocab_size); + } else { + for (int32_t i = 0; i < cfg.codec_vocab_size; ++i) { + logits[i] /= temperature; + } + + if (top_k > 0 && top_k < cfg.codec_vocab_size) { + std::vector> scored(cfg.codec_vocab_size); + for (int32_t i = 0; i < cfg.codec_vocab_size; ++i) { + scored[i] = {logits[i], i}; + } + std::partial_sort(scored.begin(), scored.begin() + top_k, scored.end(), + [](const std::pair & a, const std::pair & b) { + return a.first > b.first; + }); + float threshold = scored[top_k - 1].first; + for (int32_t i = 0; i < cfg.codec_vocab_size; ++i) { + if (logits[i] < threshold) { + logits[i] = -INFINITY; + } + } + } + + float max_logit = *std::max_element(logits.data(), logits.data() + cfg.codec_vocab_size); + double sum = 0.0; + for (int32_t i = 0; i < cfg.codec_vocab_size; ++i) { + probs[i] = expf(logits[i] - max_logit); + sum += probs[i]; + } + for (int32_t i = 0; i < cfg.codec_vocab_size; ++i) { + probs[i] = (float)(probs[i] / sum); + } + + std::discrete_distribution dist(probs.begin(), probs.end()); + next_token = dist(rng_); + } + + if (next_token == cfg.codec_eos_id) { + break; + } + + frame_codes[0] = next_token; + generated_cb0_tokens.insert(next_token); + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + std::vector codes_1_15; + if (!predict_codes_autoregressive(last_hidden_.data(), frame_codes[0], codes_1_15, temperature, top_k)) { + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + timing.t_code_pred_ms += std::chrono::duration(t1 - t0).count(); +#endif + + for (int cb = 1; cb < cfg.n_codebooks; ++cb) { + frame_codes[cb] = codes_1_15[cb - 1]; + } + + for (int cb = 0; cb < cfg.n_codebooks; ++cb) { + output.push_back(frame_codes[cb]); + } + +#ifdef QWEN3_TTS_TIMING + timing.n_frames = frame + 1; +#endif + + if (frame + 1 >= max_len) { + break; + } + + std::fill(step_embd.begin(), step_embd.end(), 0.0f); + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (!lookup_single_embedding_row(model_.codec_embd, frame_codes[0], embd_row.data())) { + return false; + } + for (int32_t h = 0; h < cfg.hidden_size; ++h) { + step_embd[h] = embd_row[h]; + } + + for (int cb = 1; cb < cfg.n_codebooks; ++cb) { + int32_t code_token = frame_codes[cb]; + if (!lookup_single_embedding_row(model_.code_pred_embd[cb - 1], code_token, embd_row.data())) { + return false; + } + for (int32_t h = 0; h < cfg.hidden_size; ++h) { + step_embd[h] += embd_row[h]; + } + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + timing.t_embed_lookup_ms += std::chrono::duration(t1 - t0).count(); +#endif + + const float * trailing_row = (frame < trailing_len) + ? trailing_text_hidden.data() + (size_t)frame * cfg.hidden_size + : tts_pad_embed.data(); + for (int32_t h = 0; h < cfg.hidden_size; ++h) { + step_embd[h] += trailing_row[h]; + } + +#ifdef QWEN3_TTS_TIMING + t0 = clk::now(); +#endif + if (!forward_step(step_embd.data(), n_past, logits)) { + return false; + } +#ifdef QWEN3_TTS_TIMING + t1 = clk::now(); + timing.t_talker_forward_ms += std::chrono::duration(t1 - t0).count(); +#endif + + n_past++; + } + +#ifdef QWEN3_TTS_TIMING + timing.t_generate_total_ms = std::chrono::duration(clk::now() - t_gen_start).count(); + timing_ = nullptr; + const auto & t = timing; + int nf = t.n_frames; + fprintf(stderr, "\n=== Detailed Generation Timing (%d frames) ===\n", nf); + fprintf(stderr, "\n Prefill:\n"); + fprintf(stderr, " Build graph: %8.1f ms\n", t.t_prefill_build_ms); + fprintf(stderr, " Forward total: %8.1f ms\n", t.t_prefill_forward_ms); + fprintf(stderr, " Graph build: %8.1f ms\n", t.t_prefill_graph_build_ms); + fprintf(stderr, " Graph alloc: %8.1f ms\n", t.t_prefill_graph_alloc_ms); + fprintf(stderr, " Compute: %8.1f ms\n", t.t_prefill_compute_ms); + fprintf(stderr, " Data I/O: %8.1f ms\n", t.t_prefill_data_ms); + fprintf(stderr, "\n Talker forward_step (total / per-frame):\n"); + fprintf(stderr, " Total: %8.1f ms (%.1f ms/frame)\n", t.t_talker_forward_ms, nf > 0 ? t.t_talker_forward_ms / nf : 0.0); + fprintf(stderr, " Graph build: %8.1f ms (%.1f ms/frame)\n", t.t_talker_graph_build_ms, nf > 0 ? t.t_talker_graph_build_ms / nf : 0.0); + fprintf(stderr, " Graph alloc: %8.1f ms (%.1f ms/frame)\n", t.t_talker_graph_alloc_ms, nf > 0 ? t.t_talker_graph_alloc_ms / nf : 0.0); + fprintf(stderr, " Compute: %8.1f ms (%.1f ms/frame)\n", t.t_talker_compute_ms, nf > 0 ? t.t_talker_compute_ms / nf : 0.0); + fprintf(stderr, " Data I/O: %8.1f ms (%.1f ms/frame)\n", t.t_talker_data_ms, nf > 0 ? t.t_talker_data_ms / nf : 0.0); + fprintf(stderr, "\n Code predictor (total / per-frame):\n"); + fprintf(stderr, " Backend: %s\n", use_coreml_code_predictor_ ? "CoreML (CPU+NE)" : "GGML"); + if (use_coreml_code_predictor_ && !coreml_code_predictor_path_.empty()) { + fprintf(stderr, " CoreML model: %s\n", coreml_code_predictor_path_.c_str()); + } + fprintf(stderr, " Total: %8.1f ms (%.1f ms/frame)\n", t.t_code_pred_ms, nf > 0 ? t.t_code_pred_ms / nf : 0.0); + fprintf(stderr, " Init/KV/embed: %8.1f ms (%.1f ms/frame)\n", t.t_code_pred_init_ms, nf > 0 ? t.t_code_pred_init_ms / nf : 0.0); + fprintf(stderr, " Prefill (2tok): %8.1f ms (%.1f ms/frame)\n", t.t_code_pred_prefill_ms, nf > 0 ? t.t_code_pred_prefill_ms / nf : 0.0); + fprintf(stderr, " Steps (14): %8.1f ms (%.1f ms/frame)\n", t.t_code_pred_steps_ms, nf > 0 ? t.t_code_pred_steps_ms / nf : 0.0); + fprintf(stderr, " Graph build: %8.1f ms (%.1f ms/frame)\n", t.t_code_pred_graph_build_ms, nf > 0 ? t.t_code_pred_graph_build_ms / nf : 0.0); + fprintf(stderr, " Graph alloc: %8.1f ms (%.1f ms/frame)\n", t.t_code_pred_graph_alloc_ms, nf > 0 ? t.t_code_pred_graph_alloc_ms / nf : 0.0); + fprintf(stderr, " Compute: %8.1f ms (%.1f ms/frame)\n", t.t_code_pred_compute_ms, nf > 0 ? t.t_code_pred_compute_ms / nf : 0.0); + fprintf(stderr, " Data I/O: %8.1f ms (%.1f ms/frame)\n", t.t_code_pred_data_ms, nf > 0 ? t.t_code_pred_data_ms / nf : 0.0); + fprintf(stderr, " CoreML total: %8.1f ms (%.1f ms/frame)\n", t.t_code_pred_coreml_ms, nf > 0 ? t.t_code_pred_coreml_ms / nf : 0.0); + fprintf(stderr, "\n Embed lookups: %8.1f ms (%.1f ms/frame)\n", t.t_embed_lookup_ms, nf > 0 ? t.t_embed_lookup_ms / nf : 0.0); + double accounted = t.t_prefill_build_ms + t.t_prefill_forward_ms + t.t_talker_forward_ms + t.t_code_pred_ms + t.t_embed_lookup_ms; + fprintf(stderr, " Other/overhead: %8.1f ms\n", t.t_generate_total_ms - accounted); + fprintf(stderr, " ─────────────────────────────────────────\n"); + fprintf(stderr, " Total generate: %8.1f ms\n", t.t_generate_total_ms); + if (nf > 0) { + fprintf(stderr, " Throughput: %8.1f ms/frame (%.1f frames/s)\n", + t.t_generate_total_ms / nf, 1000.0 * nf / t.t_generate_total_ms); + } +#endif + + return true; +} + +bool TTSTransformer::forward(const int32_t * tokens, int32_t n_tokens, int32_t n_past, + std::vector & output) { + return forward_text(tokens, n_tokens, nullptr, n_past, output); +} + +bool TTSTransformer::forward_with_audio(const int32_t * tokens, int32_t n_tokens, + const float * audio_embd, int32_t n_audio, + int32_t audio_start_pos, int32_t n_past, + std::vector & output) { + (void)audio_embd; + (void)n_audio; + (void)audio_start_pos; + return forward_text(tokens, n_tokens, nullptr, n_past, output); +} + +void free_transformer_model(tts_transformer_model & model) { + if (model.buffer) { + ggml_backend_buffer_free(model.buffer); + model.buffer = nullptr; + } + if (model.ctx) { + ggml_free(model.ctx); + model.ctx = nullptr; + } + model.tensors.clear(); + model.layers.clear(); + model.code_pred_layers.clear(); + model.code_pred_embd.clear(); + model.code_pred_head.clear(); +} + +void free_tts_kv_cache(tts_kv_cache & cache) { + if (cache.buffer) { + ggml_backend_buffer_free(cache.buffer); + cache.buffer = nullptr; + } + if (cache.ctx) { + ggml_free(cache.ctx); + cache.ctx = nullptr; + } + cache.k_cache.clear(); + cache.v_cache.clear(); + cache.n_ctx = 0; + cache.n_used = 0; +} + +} // namespace qwen3_tts diff --git a/otherarch/qwen3tts/tts_transformer.h b/otherarch/qwen3tts/tts_transformer.h new file mode 100644 index 000000000..3b228daa7 --- /dev/null +++ b/otherarch/qwen3tts/tts_transformer.h @@ -0,0 +1,353 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" +#include "gguf.h" +#include "coreml_code_predictor.h" + +#include +#include +#include +#include +#include +#ifdef QWEN3_TTS_TIMING +#include +#endif + +namespace qwen3_tts { + +#ifdef QWEN3_TTS_TIMING +struct tts_timing { + // Prefill phase + double t_prefill_build_ms = 0; // build_prefill_graph (embedding lookups, text projection) + double t_prefill_forward_ms = 0; // forward_prefill total + double t_prefill_graph_build_ms = 0; // build_prefill_forward_graph + double t_prefill_graph_alloc_ms = 0; // sched_alloc_graph + double t_prefill_compute_ms = 0; // sched_graph_compute + double t_prefill_data_ms = 0; // tensor_set + tensor_get + reset + + // Talker forward_step totals (accumulated across all frames) + double t_talker_forward_ms = 0; // total time in forward_step() + double t_talker_graph_build_ms = 0; // build_step_graph + double t_talker_graph_alloc_ms = 0; // sched_alloc_graph + double t_talker_compute_ms = 0; // sched_graph_compute + double t_talker_data_ms = 0; // tensor_set + tensor_get + reset + + // Code predictor totals (accumulated across all frames) + double t_code_pred_ms = 0; // total predict_codes_autoregressive + double t_code_pred_init_ms = 0; // init/clear KV cache + CB0 embed lookup + double t_code_pred_prefill_ms = 0; // code pred prefill (2-token, per frame) + double t_code_pred_steps_ms = 0; // code pred autoregressive steps (14 steps, per frame) + double t_code_pred_graph_build_ms = 0; // graph build (prefill + steps combined) + double t_code_pred_graph_alloc_ms = 0; // sched_alloc_graph + double t_code_pred_compute_ms = 0; // sched_graph_compute + double t_code_pred_data_ms = 0; // tensor_set + tensor_get + reset + double t_code_pred_coreml_ms = 0; // CoreML predictor compute + I/O + + // Embed lookups in generate() loop + double t_embed_lookup_ms = 0; + + int32_t n_frames = 0; + double t_generate_total_ms = 0; +}; +#endif + +#define QWEN3_TTS_MAX_NODES 16384 + +// TTS Transformer configuration (Qwen2-based Talker) +struct tts_transformer_config { + // Text embedding + int32_t text_vocab_size = 151936; + int32_t text_embd_dim = 2048; + + // Talker transformer + int32_t hidden_size = 1024; + int32_t n_layers = 28; + int32_t n_attention_heads = 16; + int32_t n_key_value_heads = 8; + int32_t intermediate_size = 3072; + int32_t head_dim = 128; + float rms_norm_eps = 1e-6f; + float rope_theta = 1000000.0f; + + // M-RoPE sections [time, freq, channel] = [24, 20, 20] + int32_t mrope_section[3] = {24, 20, 20}; + + // Codec vocabulary + int32_t codec_vocab_size = 3072; // talker.codec_embd/codec_head + int32_t n_codebooks = 16; + + // Code predictor + int32_t code_pred_layers = 5; + int32_t code_pred_vocab_size = 2048; // Per-codebook vocab + + // Special codec tokens + int32_t codec_pad_id = 2148; + int32_t codec_bos_id = 2149; + int32_t codec_eos_id = 2150; + + int32_t tts_bos_token_id = 151672; + int32_t tts_eos_token_id = 151673; + int32_t tts_pad_token_id = 151671; + + int32_t codec_think_id = 2154; + int32_t codec_nothink_id = 2155; + int32_t codec_think_bos_id = 2156; + int32_t codec_think_eos_id = 2157; + + int32_t english_language_id = 2050; +}; + +// Transformer layer weights +struct transformer_layer { + struct ggml_tensor * attn_norm = nullptr; + + struct ggml_tensor * attn_q = nullptr; + struct ggml_tensor * attn_k = nullptr; + struct ggml_tensor * attn_v = nullptr; + struct ggml_tensor * attn_output = nullptr; + struct ggml_tensor * attn_q_norm = nullptr; + struct ggml_tensor * attn_k_norm = nullptr; + + struct ggml_tensor * ffn_norm = nullptr; + + struct ggml_tensor * ffn_gate = nullptr; + struct ggml_tensor * ffn_up = nullptr; + struct ggml_tensor * ffn_down = nullptr; +}; + +// TTS Transformer model weights +struct tts_transformer_model { + tts_transformer_config config; + + // Text embedding and projection + struct ggml_tensor * text_embd = nullptr; // [text_embd_dim, text_vocab_size] + struct ggml_tensor * text_proj_fc1 = nullptr; // [text_embd_dim, text_embd_dim] + struct ggml_tensor * text_proj_fc1_bias = nullptr; + struct ggml_tensor * text_proj_fc2 = nullptr; // [text_embd_dim, hidden_size] + struct ggml_tensor * text_proj_fc2_bias = nullptr; + + // Codec embedding (for autoregressive input) + struct ggml_tensor * codec_embd = nullptr; // [hidden_size, codec_vocab_size] + + // Talker transformer layers + std::vector layers; + + // Final RMSNorm + struct ggml_tensor * output_norm = nullptr; // [hidden_size] + + // Codec head (for first codebook prediction) + struct ggml_tensor * codec_head = nullptr; // [hidden_size, codec_vocab_size] + + // Code predictor layers + std::vector code_pred_layers; + + // Code predictor output norm (final RMS norm before lm_head) + struct ggml_tensor * code_pred_output_norm = nullptr; // [hidden_size] + + // Code predictor per-codebook embeddings and heads (15 codebooks, 0 uses talker output) + std::vector code_pred_embd; // [hidden_size, code_pred_vocab_size] x 15 + std::vector code_pred_head; // [hidden_size, code_pred_vocab_size] x 15 + + // GGML context for tensor metadata + struct ggml_context * ctx = nullptr; + + // Backend buffer for weights + ggml_backend_buffer_t buffer = nullptr; + + // Tensor name to tensor mapping + std::map tensors; +}; + +// KV cache for autoregressive generation +struct tts_kv_cache { + std::vector k_cache; + std::vector v_cache; + + struct ggml_context * ctx = nullptr; + ggml_backend_buffer_t buffer = nullptr; + + int32_t n_ctx = 0; + int32_t n_used = 0; + int32_t head_dim = 128; + int32_t n_kv_heads = 8; + int32_t n_layers = 28; +}; + +// TTS Transformer state +struct tts_transformer_state { + ggml_backend_t backend = nullptr; + ggml_backend_t backend_cpu = nullptr; + ggml_backend_sched_t sched = nullptr; + + std::vector compute_meta; + + tts_kv_cache cache; // Talker KV cache (28 layers) + tts_kv_cache code_pred_cache; // Code predictor KV cache (5 layers) +}; + +// TTS Transformer class +class TTSTransformer { +public: + TTSTransformer(); + ~TTSTransformer(); + + // Load model from GGUF file + bool load_model(const std::string & model_path); + + // Release all model/runtime resources + void unload_model(); + + // Initialize KV cache + bool init_kv_cache(int32_t n_ctx); + + // Clear KV cache + void clear_kv_cache(); + + // Initialize code predictor KV cache (5 layers, max 16 context) + bool init_code_pred_kv_cache(int32_t n_ctx); + + // Clear code predictor KV cache + void clear_code_pred_kv_cache(); + + // Forward pass for text tokens (prefill phase) + // text_tokens: input text token IDs [n_tokens] + // speaker_embd: speaker embedding [hidden_size] (optional, can be nullptr) + // n_past: number of tokens already in KV cache + // output: hidden states [n_tokens, hidden_size] + bool forward_text(const int32_t * text_tokens, int32_t n_tokens, + const float * speaker_embd, int32_t n_past, + std::vector & output); + + bool forward_prefill(const float * prefill_embd, int32_t n_tokens, + int32_t n_past, std::vector & output, + std::vector * logits_out = nullptr); + + // Forward pass for codec tokens (generation phase) + // codec_token: single codec token for first codebook + // n_past: number of tokens already in KV cache + // output: logits for next codec token [codec_vocab_size] + bool forward_codec(int32_t codec_token, int32_t n_past, + std::vector & output); + + bool forward_step(const float * step_embd, int32_t n_past, + std::vector & output, + std::vector * hidden_out = nullptr); + + // Get hidden states from last forward pass (for code predictor) + bool get_hidden_states(std::vector & hidden) const; + + // Run code predictor to get all 16 codebook predictions + // hidden: hidden states from talker [hidden_size] + // prev_codes: previous codes for codebooks 1-15 (can be nullptr for first step) + // output: logits for all 16 codebooks [16, code_pred_vocab_size] + bool predict_codes(const float * hidden, const int32_t * prev_codes, + std::vector & output); + + // Run code predictor autoregressively to generate 15 codes (codebooks 1-15) + // hidden: hidden states from talker [hidden_size] + // codebook_0_token: the codebook 0 token (used to create 2-token prefill input) + // output: generated codes for codebooks 1-15 [15] + bool predict_codes_autoregressive(const float * hidden, int32_t codebook_0_token, + std::vector & output, + float temperature = 0.9f, + int32_t top_k = 50); + + // Generate speech codes autoregressively + // text_tokens: input text token IDs [n_tokens] + // speaker_embd: speaker embedding [hidden_size] + // max_len: maximum number of frames to generate + // output: generated speech codes [n_frames, n_codebooks] + bool generate(const int32_t * text_tokens, int32_t n_tokens, + const float * speaker_embd, int32_t max_len, + std::vector & output, + int32_t language_id = 2050, + float repetition_penalty = 1.05f, + float temperature = 0.9f, + int32_t top_k = 50); + + const tts_transformer_config & get_config() const { return model_.config; } + + const std::string & get_error() const { return error_msg_; } + + // Legacy interface for compatibility + bool forward(const int32_t * tokens, int32_t n_tokens, int32_t n_past, + std::vector & output); + + bool forward_with_audio(const int32_t * tokens, int32_t n_tokens, + const float * audio_embd, int32_t n_audio, + int32_t audio_start_pos, int32_t n_past, + std::vector & output); + +private: + bool try_init_coreml_code_predictor(const std::string & model_path); + bool predict_codes_autoregressive_coreml(const float * hidden, int32_t codebook_0_token, + std::vector & output, + float temperature, + int32_t top_k); + + bool build_prefill_graph(const int32_t * text_tokens, int32_t n_tokens, + const float * speaker_embd, int32_t language_id, + std::vector & prefill_embd, + std::vector & trailing_text_hidden, + std::vector & tts_pad_embed); + + struct ggml_cgraph * build_prefill_forward_graph(int32_t n_tokens, int32_t n_past); + + struct ggml_cgraph * build_step_graph(int32_t n_past); + + bool project_text_tokens(const int32_t * text_tokens, int32_t n_tokens, + std::vector & output); + + bool lookup_embedding_rows(struct ggml_tensor * embedding, const int32_t * token_ids, + int32_t n_tokens, const char * input_name, + const char * output_name, std::vector & output); + bool lookup_single_embedding_row(struct ggml_tensor * embedding, int32_t token_id, + float * out_row); + + // Build computation graph for code predictor + struct ggml_cgraph * build_code_pred_graph(int32_t n_prev_codes); + + // Build computation graph for single-step autoregressive code predictor + // n_past: number of tokens already in KV cache (0-14) + // generation_step: which codebook we're predicting (0-14) + struct ggml_cgraph * build_code_pred_step_graph(int32_t n_past, int32_t generation_step); + + // Build computation graph for 2-token prefill of code predictor + // Processes [past_hidden, codec_embd(codebook_0_token)] together + struct ggml_cgraph * build_code_pred_prefill_graph(); + + // Parse hyperparameters from GGUF + bool parse_config(struct gguf_context * ctx); + + // Create tensor structures + bool create_tensors(struct gguf_context * ctx); + + // Load tensor data from file + bool load_tensor_data(const std::string & path, struct gguf_context * ctx); + + tts_transformer_model model_; + tts_transformer_state state_; + std::string error_msg_; + + // Cached hidden states from last forward pass + std::vector last_hidden_; + std::vector embd_row_fp16_scratch_; + std::mt19937 rng_{std::random_device{}()}; + CoreMLCodePredictor coreml_code_predictor_; + bool use_coreml_code_predictor_ = false; + std::string coreml_code_predictor_path_; + bool skip_ggml_code_pred_layers_ = false; + +#ifdef QWEN3_TTS_TIMING + tts_timing * timing_ = nullptr; +#endif +}; + +// Free model resources +void free_transformer_model(tts_transformer_model & model); + +// Free KV cache resources +void free_tts_kv_cache(tts_kv_cache & cache); + +} // namespace qwen3_tts