From 0aa35ef05261367979f43ca5f82470dfec701b9a Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Mon, 9 Mar 2026 15:07:06 +0800 Subject: [PATCH 01/10] kokoro TTS Supported axcl backend --- sherpa-onnx/csrc/CMakeLists.txt | 5 + sherpa-onnx/csrc/axcl/axcl-model.cc | 5 + sherpa-onnx/csrc/axcl/axcl-model.h | 3 + .../axcl/offline-tts-kokoro-model-axcl.cc | 596 ++++++++++++++++++ .../csrc/axcl/offline-tts-kokoro-model-axcl.h | 37 ++ sherpa-onnx/csrc/offline-tts-impl.cc | 33 +- sherpa-onnx/csrc/offline-tts-kokoro-impl.h | 61 +- 7 files changed, 714 insertions(+), 26 deletions(-) create mode 100644 sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc create mode 100644 sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h diff --git a/sherpa-onnx/csrc/CMakeLists.txt b/sherpa-onnx/csrc/CMakeLists.txt index 1a34d1c1b0..52b106c942 100755 --- a/sherpa-onnx/csrc/CMakeLists.txt +++ b/sherpa-onnx/csrc/CMakeLists.txt @@ -223,6 +223,11 @@ if(SHERPA_ONNX_ENABLE_AXCL) ./axcl/offline-sense-voice-model-axcl.cc ./axcl/utils.cc ) + if(SHERPA_ONNX_ENABLE_TTS) + list(APPEND sources + ./axcl/offline-tts-kokoro-model-axcl.cc + ) + endif() endif() if(SHERPA_ONNX_ENABLE_RKNN OR SHERPA_ONNX_ENABLE_ASCEND_NPU OR SHERPA_ONNX_ENABLE_QNN OR SHERPA_ONNX_ENABLE_AXERA OR SHERPA_ONNX_ENABLE_AXCL) diff --git a/sherpa-onnx/csrc/axcl/axcl-model.cc b/sherpa-onnx/csrc/axcl/axcl-model.cc index 85a3b8a8f1..02fc19257a 100644 --- a/sherpa-onnx/csrc/axcl/axcl-model.cc +++ b/sherpa-onnx/csrc/axcl/axcl-model.cc @@ -429,6 +429,11 @@ bool AxclModel::SetInputTensorData(const std::string &name, const int32_t *p, return impl_->SetInputTensorData(name, p, n); } +bool AxclModel::SetInputTensorData(const std::string &name, const uint8_t *p, + int32_t n) const { + return impl_->SetInputTensorData(name, p, n); +} + std::vector AxclModel::GetOutputTensorData( const std::string &name) const { return impl_->GetOutputTensorData(name); diff --git a/sherpa-onnx/csrc/axcl/axcl-model.h b/sherpa-onnx/csrc/axcl/axcl-model.h index e6685bae0f..e43c9f2d05 100644 --- a/sherpa-onnx/csrc/axcl/axcl-model.h +++ b/sherpa-onnx/csrc/axcl/axcl-model.h @@ -33,6 +33,9 @@ class AxclModel { bool SetInputTensorData(const std::string &name, const int32_t *p, int32_t n) const; + + bool SetInputTensorData(const std::string &name, const uint8_t *p, + int32_t n) const; std::vector GetOutputTensorData(const std::string &name) const; diff --git a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc new file mode 100644 index 0000000000..fd5f65deed --- /dev/null +++ b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc @@ -0,0 +1,596 @@ +// sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc +// +// Copyright (c) 2025 Xiaomi Corporation + +#include "sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if __ANDROID_API__ >= 9 +#include "android/asset_manager.h" +#include "android/asset_manager_jni.h" +#endif + +#if __OHOS__ +#include "rawfile/raw_file_manager.h" +#endif + +#include "kaldi-native-fbank/csrc/istft.h" +#include "sherpa-onnx/csrc/axcl/axcl-model.h" +#include "sherpa-onnx/csrc/file-utils.h" +#include "sherpa-onnx/csrc/macros.h" +#include "sherpa-onnx/csrc/onnx-utils.h" +#include "sherpa-onnx/csrc/session.h" + +namespace sherpa_onnx { + +static std::vector Sigmoid(const std::vector &x) { + std::vector result(x.size()); + for (size_t i = 0; i < x.size(); i++) { + result[i] = 1.0f / (1.0f + std::exp(-x[i])); + } + return result; +} + +template +static std::vector Argsort(const std::vector &v, int32_t len, + bool reverse) { + std::vector idx(len); + std::iota(idx.begin(), idx.end(), 0); + + if (!reverse) { + std::stable_sort(idx.begin(), idx.end(), + [&v](size_t i1, size_t i2) { return v[i1] < v[i2]; }); + } else { + std::stable_sort(idx.begin(), idx.end(), + [&v](size_t i1, size_t i2) { return v[i1] > v[i2]; }); + } + + return idx; +} + +template +static std::vector NpRepeat(const std::vector &v, + const std::vector ×) { + std::vector result; + for (size_t i = 0; i < times.size(); i++) { + for (int32_t n = 0; n < times[i]; n++) { + result.push_back(v[i]); + } + } + return result; +} + +template +static std::vector Linspace(T a, T b, size_t N) { + if (N == 0) return {}; + if (N == 1) return {a}; + T h = (b - a) / static_cast(N - 1); + std::vector xs(N); + T val = a; + for (size_t i = 0; i < N; ++i) { + xs[i] = val; + val += h; + } + return xs; +} + +class OfflineTtsKokoroModelAxcl::Impl { + private: + int32_t max_seq_len_ = 0; + std::vector styles_; + + public: + explicit Impl(const OfflineTtsModelConfig &config) + : config_(config), + env_(ORT_LOGGING_LEVEL_ERROR), + sess_opts_(GetSessionOptions(config)), + allocator_{} { + Init(config); + } + + template + Impl(Manager *mgr, const OfflineTtsModelConfig &config) + : config_(config), + env_(ORT_LOGGING_LEVEL_ERROR), + sess_opts_(GetSessionOptions(config)), + allocator_{} { + Init(mgr, config); + } + + const OfflineTtsKokoroModelMetaData &GetMetaData() const { + return meta_data_; + } + + std::vector Run(std::vector input_ids, int64_t sid, float speed) { + if (input_ids.size() < 2) { + return {}; + } + + if (config_.kokoro.length_scale != 1 && speed == 1) { + speed = 1. / config_.kokoro.length_scale; + } + + int32_t phoneme_len = static_cast(input_ids.size()) - 2; + if (phoneme_len < 0) phoneme_len = 0; + + std::vector ref_s = LoadVoiceEmbedding(sid, phoneme_len); + + std::vector audio; + bool success = RunModels(input_ids, ref_s, speed, audio); + if (!success) { + SHERPA_ONNX_LOGE("Run models failed!"); + return {}; + } + + return audio; + } + + private: + void Init(const OfflineTtsModelConfig &config) { + std::string model_path = config.kokoro.model; + + // kokoro axcl mode we assume config.kokoro.model is a directory containing models + // or we can use the pattern like kokoro.model + "_part1.axmodel" + // Let's assume the user passes a directory. Or the user passes the prefix path. + // For simplicity, we assume `config.kokoro.model` is a directory. + std::string model1_path = model_path + "/kokoro_part1.axmodel"; + std::string model2_path = model_path + "/kokoro_part2.axmodel"; + std::string model3_path = model_path + "/kokoro_part3.axmodel"; + std::string model4_path = model_path + "/kokoro_part4.onnx"; + + model1_ = std::make_unique(model1_path); + model2_ = std::make_unique(model2_path); + model3_ = std::make_unique(model3_path); + + auto buf4 = ReadFile(model4_path); + model4_ = std::make_unique(env_, buf4.data(), buf4.size(), + sess_opts_); + GetInputNames(model4_.get(), &model4_input_names_, &model4_input_names_ptr_); + GetOutputNames(model4_.get(), &model4_output_names_, + &model4_output_names_ptr_); + + auto voices_buf = ReadFile(config.kokoro.voices); + LoadVoices(voices_buf.data(), voices_buf.size()); + } + + template + void Init(Manager *mgr, const OfflineTtsModelConfig &config) { + std::string model_path = config.kokoro.model; + std::string model1_path = model_path + "/kokoro_part1.axmodel"; + std::string model2_path = model_path + "/kokoro_part2.axmodel"; + std::string model3_path = model_path + "/kokoro_part3.axmodel"; + std::string model4_path = model_path + "/kokoro_part4.onnx"; + + auto buf1 = ReadFile(mgr, model1_path); + model1_ = std::make_unique(buf1.data(), buf1.size()); + + auto buf2 = ReadFile(mgr, model2_path); + model2_ = std::make_unique(buf2.data(), buf2.size()); + + auto buf3 = ReadFile(mgr, model3_path); + model3_ = std::make_unique(buf3.data(), buf3.size()); + + auto buf4 = ReadFile(mgr, model4_path); + model4_ = std::make_unique(env_, buf4.data(), buf4.size(), + sess_opts_); + GetInputNames(model4_.get(), &model4_input_names_, &model4_input_names_ptr_); + GetOutputNames(model4_.get(), &model4_output_names_, + &model4_output_names_ptr_); + + auto voices_buf = ReadFile(mgr, config.kokoro.voices); + LoadVoices(voices_buf.data(), voices_buf.size()); + } + + void LoadVoices(const char *voices_data, size_t voices_data_length) { + // get metadata + meta_data_.sample_rate = 24000; + meta_data_.num_speakers = voices_data_length / (sizeof(float) * 256 * 510); + meta_data_.version = 1; + meta_data_.has_espeak = 1; + meta_data_.voice = "en-us"; + meta_data_.max_token_len = 510; + + int32_t actual_num_floats = voices_data_length / sizeof(float); + styles_ = std::vector( + reinterpret_cast(voices_data), + reinterpret_cast(voices_data) + actual_num_floats); + + auto shape1 = model1_->TensorShape("input_ids"); + max_seq_len_ = shape1[1]; + } + + std::vector LoadVoiceEmbedding(int32_t sid, int32_t phoneme_len) { + int32_t style_dim0 = 510; // max_token_len + int32_t style_dim1 = 256; + phoneme_len = std::max(phoneme_len, 0); + + std::vector ref_s(style_dim1); + if (phoneme_len < style_dim0) { + const float *p = styles_.data() + sid * style_dim0 * style_dim1 + + phoneme_len * style_dim1; + std::copy(p, p + style_dim1, ref_s.begin()); + } else { + int32_t idx = style_dim0 / 2; + const float *p = + styles_.data() + sid * style_dim0 * style_dim1 + idx * style_dim1; + std::copy(p, p + style_dim1, ref_s.begin()); + } + return ref_s; + } + + bool RunModels(std::vector &input_ids, + const std::vector &ref_s, float speed, + std::vector &audio) { + int32_t actual_len = input_ids.size(); + if (actual_len > max_seq_len_) { + SHERPA_ONNX_LOGE("Input ids length %d exceeds max_seq_len %d, truncating.", + actual_len, max_seq_len_); + input_ids.resize(max_seq_len_); + actual_len = max_seq_len_; + } + + int32_t padding_len = max_seq_len_ - actual_len; + if (padding_len > 0) { + std::vector padding(padding_len, 0); + input_ids.insert(input_ids.end(), padding.begin(), padding.end()); + } + + int32_t actual_content_frames = 0; + int32_t total_frames = 0; + if (!InferenceSingleChunk(input_ids, ref_s, actual_len, speed, audio, + actual_content_frames, total_frames)) { + return false; + } + + TrimAudioByContent(audio, actual_content_frames, total_frames, actual_len); + + return true; + } + + bool InferenceSingleChunk(std::vector &input_ids, + const std::vector &ref_s, int32_t actual_len, + float speed, std::vector &audio, + int32_t &actual_content_frames, + int32_t &total_frames) { + bool is_doubled = false; + int32_t original_actual_len = actual_len; + + PrepareInputIds(input_ids, actual_len, is_doubled); + + std::vector input_ids_i32(input_ids.begin(), input_ids.end()); + + std::vector input_lengths; + std::vector text_mask; + ComputeExternalPreprocessing(input_ids, actual_len, input_lengths, + text_mask); + + std::vector text_mask_u8(text_mask.begin(), text_mask.end()); + + model1_->SetInputTensorData("input_ids", input_ids_i32.data(), + input_ids_i32.size()); + model1_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); + model1_->SetInputTensorData("text_mask", text_mask_u8.data(), text_mask_u8.size()); + + if (!model1_->Run()) { + return false; + } + + std::vector duration = model1_->GetOutputTensorData("duration"); + std::vector d = model1_->GetOutputTensorData("d"); + + std::vector pred_dur; + ProcessDuration(duration, actual_len, speed, pred_dur, total_frames); + + std::vector pred_aln_trg = CreateAlignmentMatrix(pred_dur, total_frames); + + int32_t d_shape_1 = max_seq_len_; + int32_t d_shape_2 = 640; // style_dim_[2] essentially, or based on d_shape_[2] + std::vector en(total_frames * d_shape_2, 0.0f); + + for (int32_t j = 0; j < d_shape_2; ++j) { + for (int32_t k = 0; k < total_frames; ++k) { + float sum = 0.0f; + for (int32_t i = 0; i < max_seq_len_; ++i) { + sum += d[i * d_shape_2 + j] * pred_aln_trg[i * total_frames + k]; + } + en[j * total_frames + k] = sum; + } + } + + std::vector text_mask_float(text_mask.begin(), text_mask.end()); + + model2_->SetInputTensorData("en", en.data(), en.size()); + model2_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); + model2_->SetInputTensorData("input_ids", input_ids_i32.data(), input_ids_i32.size()); + model2_->SetInputTensorData("text_mask", text_mask_float.data(), text_mask_float.size()); + model2_->SetInputTensorData("pred_aln_trg", pred_aln_trg.data(), pred_aln_trg.size()); + + if (!model2_->Run()) { + return false; + } + + std::vector f0_pred = model2_->GetOutputTensorData("F0_pred"); + std::vector n_pred = model2_->GetOutputTensorData("N_pred"); + std::vector asr = model2_->GetOutputTensorData("asr"); + + auto memory_info = + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + std::array f0_shape = {1, static_cast(f0_pred.size())}; + Ort::Value f0_tensor = Ort::Value::CreateTensor( + memory_info, f0_pred.data(), f0_pred.size(), f0_shape.data(), + f0_shape.size()); + + auto out4 = model4_->Run({}, model4_input_names_ptr_.data(), &f0_tensor, 1, + model4_output_names_ptr_.data(), + model4_output_names_ptr_.size()); + float *p_har = out4[0].GetTensorMutableData(); + auto har_shape = out4[0].GetTensorTypeAndShapeInfo().GetShape(); + int64_t har_num_elements = 1; + for (auto d : har_shape) { + har_num_elements *= d; + } + std::vector har(p_har, p_har + har_num_elements); + + model3_->SetInputTensorData("asr", asr.data(), asr.size()); + model3_->SetInputTensorData("F0_pred", f0_pred.data(), f0_pred.size()); + model3_->SetInputTensorData("N_pred", n_pred.data(), n_pred.size()); + model3_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); + model3_->SetInputTensorData("har", har.data(), har.size()); + + if (!model3_->Run()) { + return false; + } + + std::vector x_out = model3_->GetOutputTensorData("x"); + PostprocessXToAudio(x_out, 23041, audio); + + if (is_doubled) { + actual_content_frames = std::accumulate(pred_dur.begin(), + pred_dur.begin() + original_actual_len, 0); + audio.erase(audio.begin() + audio.size() / 2, audio.end()); + total_frames = total_frames / 2; + } else { + actual_content_frames = + std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); + } + return true; + } + + void ProcessDuration(const std::vector &duration, int32_t actual_len, + float speed, std::vector &pred_dur, + int32_t &total_frames) { + std::vector pred_dur_original(actual_len, 0); + std::vector duration_processed = Sigmoid(duration); + + int32_t dur_shape_2 = 50; + for (int32_t i = 0; i < actual_len; i++) { + float sum = 0; + for (int32_t n = 0; n < dur_shape_2; n++) { + sum += duration_processed[i * dur_shape_2 + n]; + } + sum /= speed; + pred_dur_original[i] = static_cast(std::max(1.f, std::round(sum))); + } + + std::vector pred_dur_padding(max_seq_len_ - actual_len, 0); + pred_dur = pred_dur_original; + pred_dur.insert(pred_dur.end(), pred_dur_padding.begin(), + pred_dur_padding.end()); + + int32_t fixed_total_frames = max_seq_len_ * 2; + int32_t actual_frames = + std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); + int32_t diff = fixed_total_frames - actual_frames; + + if (diff < 0) { + auto indices = Argsort(pred_dur, actual_len, true); + int32_t decreased = 0; + for (auto idx : indices) { + if (pred_dur[idx] > 1 && decreased < std::abs(diff)) { + pred_dur[idx]--; + decreased++; + } + if (decreased >= std::abs(diff)) break; + } + } + + actual_frames = + std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); + int32_t remaining_frames = fixed_total_frames - actual_frames; + int32_t padding_len = max_seq_len_ - actual_len; + + if (remaining_frames > 0 && padding_len > 0) { + int32_t frames_per_padding = remaining_frames / padding_len; + int32_t remainder = remaining_frames % padding_len; + + for (int32_t i = actual_len; i < static_cast(pred_dur.size()); i++) + pred_dur[i] = frames_per_padding; + + if (remainder > 0) { + for (int32_t i = actual_len; i < actual_len + remainder; i++) + pred_dur[i] += 1; + } + } + + total_frames = std::accumulate(pred_dur.begin(), pred_dur.end(), 0); + } + + std::vector CreateAlignmentMatrix(const std::vector &pred_dur, + int32_t total_frames) { + std::vector seq_range(max_seq_len_); + std::iota(seq_range.begin(), seq_range.end(), 0); + auto indices = NpRepeat(seq_range, pred_dur); + + std::vector pred_aln_trg(max_seq_len_ * total_frames, 0.0f); + if (!indices.empty()) { + int32_t col = 0; + for (auto i : indices) { + pred_aln_trg[i * total_frames + col] = 1.0f; + col++; + } + } + + return pred_aln_trg; + } + + void PostprocessXToAudio(const std::vector &x, int32_t num_frames, + std::vector &audio) { + int32_t n_fft = 20; + int32_t hop_length = 5; + int32_t half_n_fft = n_fft / 2 + 1; + + std::vector spec_part(half_n_fft * num_frames); + std::vector phase_part(half_n_fft * num_frames); + std::vector cos_part(half_n_fft * num_frames); + spec_part.assign(x.begin(), x.begin() + half_n_fft * num_frames); + phase_part.assign(x.begin() + half_n_fft * num_frames, x.end()); + + for (int32_t i = 0; i < half_n_fft * num_frames; i++) { + spec_part[i] = std::exp(spec_part[i]); + phase_part[i] = std::sin(phase_part[i]); + cos_part[i] = + std::sqrt(1.f - std::max(0.f, std::min(phase_part[i] * phase_part[i], 1.0f))); + } + + knf::StftResult stft_result; + stft_result.num_frames = num_frames; + stft_result.real.resize(half_n_fft * num_frames); + stft_result.imag.resize(half_n_fft * num_frames); + + for (int32_t i = 0; i < half_n_fft; i++) { + for (int32_t n = 0; n < num_frames; n++) { + float spec = spec_part[i * num_frames + n]; + float real_part = spec * cos_part[i * num_frames + n]; + float imag_part = spec * phase_part[i * num_frames + n]; + + // StftResult layout is [frame_index * half_n_fft + bin] + stft_result.real[n * half_n_fft + i] = real_part; + stft_result.imag[n * half_n_fft + i] = imag_part; + } + } + + knf::StftConfig stft_config; + stft_config.n_fft = n_fft; + stft_config.hop_length = hop_length; + stft_config.win_length = n_fft; + stft_config.center = true; + stft_config.normalized = false; + stft_config.window_type = "hann"; + stft_config.pad_mode = "reflect"; + + knf::IStft istft(stft_config); + audio = istft.Compute(stft_result); + } + + void TrimAudioByContent(std::vector &audio, int32_t actual_content_frames, + int32_t total_frames, int32_t actual_len) { + int32_t padding_len = max_seq_len_ - actual_len; + if (padding_len > 0) { + float content_ratio = actual_content_frames * 1.0f / total_frames; + int32_t audio_len_to_keep = static_cast(audio.size() * content_ratio); + audio.resize(audio_len_to_keep); + } + } + + void PrepareInputIds(std::vector &input_ids, int32_t &actual_len, + bool &is_doubled) { + is_doubled = false; + int32_t original_actual_len = actual_len; + + if (actual_len <= 32) { // DOUBLE_INPUT_THRESHOLD + is_doubled = true; + std::vector valid_content(input_ids.begin(), + input_ids.begin() + actual_len); + std::vector input_ids_doubled; + input_ids_doubled.insert(input_ids_doubled.end(), valid_content.begin(), + valid_content.end()); + input_ids_doubled.insert(input_ids_doubled.end(), valid_content.begin(), + valid_content.end()); + + int32_t padding_len = max_seq_len_ - 2 * actual_len; + if (padding_len > 0) { + std::vector padding(padding_len, 0); + input_ids_doubled.insert(input_ids_doubled.end(), padding.begin(), + padding.end()); + } else { + input_ids_doubled.resize(max_seq_len_); + } + + input_ids = input_ids_doubled; + actual_len = std::min(original_actual_len * 2, max_seq_len_); + } + } + + void ComputeExternalPreprocessing(const std::vector &input_ids, + int32_t actual_len, + std::vector &input_lengths, + std::vector &text_mask) { + input_lengths = std::vector{actual_len}; + text_mask.resize(max_seq_len_); + for (int32_t i = 0; i < max_seq_len_; i++) { + text_mask[i] = (i >= actual_len) ? 1 : 0; + } + } + + private: + OfflineTtsModelConfig config_; + OfflineTtsKokoroModelMetaData meta_data_; + + Ort::Env env_; + Ort::SessionOptions sess_opts_; + Ort::AllocatorWithDefaultOptions allocator_; + + std::unique_ptr model1_; + std::unique_ptr model2_; + std::unique_ptr model3_; + + std::unique_ptr model4_; + std::vector model4_input_names_; + std::vector model4_input_names_ptr_; + std::vector model4_output_names_; + std::vector model4_output_names_ptr_; +}; + +OfflineTtsKokoroModelAxcl::~OfflineTtsKokoroModelAxcl() = default; + +OfflineTtsKokoroModelAxcl::OfflineTtsKokoroModelAxcl( + const OfflineTtsModelConfig &config) + : impl_(std::make_unique(config)) {} + +template +OfflineTtsKokoroModelAxcl::OfflineTtsKokoroModelAxcl( + Manager *mgr, const OfflineTtsModelConfig &config) + : impl_(std::make_unique(mgr, config)) {} + +std::vector OfflineTtsKokoroModelAxcl::Run( + const std::vector& x, int64_t sid, float speed) const { + return impl_->Run(x, sid, speed); +} + +const OfflineTtsKokoroModelMetaData &OfflineTtsKokoroModelAxcl::GetMetaData() + const { + return impl_->GetMetaData(); +} + +#if __ANDROID_API__ >= 9 +template OfflineTtsKokoroModelAxcl::OfflineTtsKokoroModelAxcl( + AAssetManager *mgr, const OfflineTtsModelConfig &config); +#endif + +#if __OHOS__ +template OfflineTtsKokoroModelAxcl::OfflineTtsKokoroModelAxcl( + NativeResourceManager *mgr, const OfflineTtsModelConfig &config); +#endif + +} // namespace sherpa_onnx + diff --git a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h new file mode 100644 index 0000000000..0430aaf687 --- /dev/null +++ b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h @@ -0,0 +1,37 @@ +// sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h +// +// Copyright (c) 2025 Xiaomi Corporation + +#ifndef SHERPA_ONNX_CSRC_AXCL_OFFLINE_TTS_KOKORO_MODEL_AXCL_H_ +#define SHERPA_ONNX_CSRC_AXCL_OFFLINE_TTS_KOKORO_MODEL_AXCL_H_ + +#include +#include + +#include "sherpa-onnx/csrc/offline-tts-kokoro-model-meta-data.h" +#include "sherpa-onnx/csrc/offline-tts-model-config.h" + +namespace sherpa_onnx { + +class OfflineTtsKokoroModelAxcl { + public: + ~OfflineTtsKokoroModelAxcl(); + + explicit OfflineTtsKokoroModelAxcl(const OfflineTtsModelConfig &config); + + template + OfflineTtsKokoroModelAxcl(Manager *mgr, const OfflineTtsModelConfig &config); + + std::vector Run(const std::vector& x, int64_t sid = 0, float speed = 1.0) const; + + const OfflineTtsKokoroModelMetaData &GetMetaData() const; + + private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace sherpa_onnx + +#endif // SHERPA_ONNX_CSRC_AXCL_OFFLINE_TTS_KOKORO_MODEL_AXCL_H_ + diff --git a/sherpa-onnx/csrc/offline-tts-impl.cc b/sherpa-onnx/csrc/offline-tts-impl.cc index 7720a49113..de1a0173da 100644 --- a/sherpa-onnx/csrc/offline-tts-impl.cc +++ b/sherpa-onnx/csrc/offline-tts-impl.cc @@ -24,6 +24,10 @@ #include "sherpa-onnx/csrc/offline-tts-vits-impl.h" #include "sherpa-onnx/csrc/offline-tts-zipvoice-impl.h" +#ifdef SHERPA_ONNX_ENABLE_AXCL +#include "sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h" +#endif + namespace sherpa_onnx { std::vector OfflineTtsImpl::AddBlank(const std::vector &x, @@ -40,6 +44,29 @@ std::vector OfflineTtsImpl::AddBlank(const std::vector &x, std::unique_ptr OfflineTtsImpl::Create( const OfflineTtsConfig &config) { + if (config.model.provider == "axcl") { +#if SHERPA_ONNX_ENABLE_AXCL + if (!config.model.kokoro.model.empty()) { + return std::make_unique>( + config); + } else { + SHERPA_ONNX_LOGE( + "Only Kokoro models are currently supported by axcl for " + "non-streaming TTS."); + SHERPA_ONNX_EXIT(-1); + return nullptr; + } + +#else + SHERPA_ONNX_LOGE( + "Please rebuild sherpa-onnx with -DSHERPA_ONNX_ENABLE_AXCL=ON if you " + "want to use axcl. See also " + "https://k2-fsa.github.io/sherpa/onnx/axcl/install.html"); + SHERPA_ONNX_EXIT(-1); + return nullptr; +#endif + } + if (!config.model.vits.model.empty()) { return std::make_unique(config); } else if (!config.model.matcha.acoustic_model.empty()) { @@ -48,7 +75,8 @@ std::unique_ptr OfflineTtsImpl::Create( !config.model.zipvoice.decoder.empty()) { return std::make_unique(config); } else if (!config.model.kokoro.model.empty()) { - return std::make_unique(config); + return std::make_unique>( + config); } else if (!config.model.kitten.model.empty()) { return std::make_unique(config); } else if (!config.model.pocket.lm_flow.empty()) { @@ -73,7 +101,8 @@ std::unique_ptr OfflineTtsImpl::Create( !config.model.zipvoice.decoder.empty()) { return std::make_unique(mgr, config); } else if (!config.model.kokoro.model.empty()) { - return std::make_unique(mgr, config); + return std::make_unique>( + mgr, config); } else if (!config.model.kitten.model.empty()) { return std::make_unique(mgr, config); } else if (!config.model.pocket.lm_flow.empty()) { diff --git a/sherpa-onnx/csrc/offline-tts-kokoro-impl.h b/sherpa-onnx/csrc/offline-tts-kokoro-impl.h index 0ca2921114..6febb001cb 100644 --- a/sherpa-onnx/csrc/offline-tts-kokoro-impl.h +++ b/sherpa-onnx/csrc/offline-tts-kokoro-impl.h @@ -7,8 +7,9 @@ #include #include #include -#include #include +#include +#include #include #include @@ -25,13 +26,17 @@ #include "sherpa-onnx/csrc/piper-phonemize-lexicon.h" #include "sherpa-onnx/csrc/text-utils.h" +#ifdef SHERPA_ONNX_ENABLE_AXCL +#include "sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h" +#endif + namespace sherpa_onnx { +template class OfflineTtsKokoroImpl : public OfflineTtsImpl { public: explicit OfflineTtsKokoroImpl(const OfflineTtsConfig &config) - : config_(config), - model_(std::make_unique(config.model)) { + : config_(config), model_(std::make_unique(config.model)) { InitFrontend(); if (!config.rule_fsts.empty()) { @@ -87,7 +92,7 @@ class OfflineTtsKokoroImpl : public OfflineTtsImpl { template OfflineTtsKokoroImpl(Manager *mgr, const OfflineTtsConfig &config) : config_(config), - model_(std::make_unique(mgr, config.model)) { + model_(std::make_unique(mgr, config.model)) { InitFrontend(mgr); if (!config.rule_fsts.empty()) { @@ -137,8 +142,8 @@ class OfflineTtsKokoroImpl : public OfflineTtsImpl { tn_list_.push_back( std::make_unique(std::move(r))); } // for (; !reader->Done(); reader->Next()) - } // for (const auto &f : files) - } // if (!config.rule_fars.empty()) + } // for (const auto &f : files) + } // if (!config.rule_fars.empty()) } int32_t SampleRate() const override { @@ -394,29 +399,37 @@ class OfflineTtsKokoroImpl : public OfflineTtsImpl { x.insert(x.end(), k.begin(), k.end()); } - auto memory_info = - Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + GeneratedAudio ans; + ans.sample_rate = model_->GetMetaData().sample_rate; - std::array x_shape = {1, static_cast(x.size())}; - Ort::Value x_tensor = Ort::Value::CreateTensor( - memory_info, x.data(), x.size(), x_shape.data(), x_shape.size()); +#ifdef SHERPA_ONNX_ENABLE_AXCL + if constexpr (std::is_same_v) { + std::vector audio = model_->Run(x, sid, speed); + ans.samples = std::move(audio); + } else +#endif + { + auto memory_info = + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); - Ort::Value audio = model_->Run(std::move(x_tensor), sid, speed); + std::array x_shape = {1, static_cast(x.size())}; + Ort::Value x_tensor = Ort::Value::CreateTensor( + memory_info, x.data(), x.size(), x_shape.data(), x_shape.size()); - std::vector audio_shape = - audio.GetTensorTypeAndShapeInfo().GetShape(); + Ort::Value audio = model_->Run(std::move(x_tensor), sid, speed); - int64_t total = 1; - // The output shape may be (1, 1, total) or (1, total) or (total,) - for (auto i : audio_shape) { - total *= i; - } + std::vector audio_shape = + audio.GetTensorTypeAndShapeInfo().GetShape(); - const float *p = audio.GetTensorData(); + int64_t total = 1; + // The output shape may be (1, 1, total) or (1, total) or (total,) + for (auto i : audio_shape) { + total *= i; + } - GeneratedAudio ans; - ans.sample_rate = model_->GetMetaData().sample_rate; - ans.samples = std::vector(p, p + total); + const float *p = audio.GetTensorData(); + ans.samples = std::vector(p, p + total); + } float silence_scale = config_.silence_scale; if (silence_scale != 1) { @@ -428,7 +441,7 @@ class OfflineTtsKokoroImpl : public OfflineTtsImpl { private: OfflineTtsConfig config_; - std::unique_ptr model_; + std::unique_ptr model_; std::vector> tn_list_; std::unique_ptr frontend_; }; From 04631218a9c7be6594e4361fee5942bb748a27fa Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Mon, 9 Mar 2026 16:11:35 +0800 Subject: [PATCH 02/10] fix offline tts kokoro axcl impl --- .../csrc/axcl/offline-tts-kokoro-axcl-impl.h | 450 ++++++++++++++++++ sherpa-onnx/csrc/offline-tts-impl.cc | 10 +- sherpa-onnx/csrc/offline-tts-kokoro-impl.h | 61 +-- 3 files changed, 478 insertions(+), 43 deletions(-) create mode 100644 sherpa-onnx/csrc/axcl/offline-tts-kokoro-axcl-impl.h diff --git a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-axcl-impl.h b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-axcl-impl.h new file mode 100644 index 0000000000..3fccb7e044 --- /dev/null +++ b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-axcl-impl.h @@ -0,0 +1,450 @@ +// sherpa-onnx/csrc/axcl/offline-tts-kokoro-axcl-impl.h +// +// Copyright (c) 2025 Xiaomi Corporation +#ifndef SHERPA_ONNX_CSRC_AXCL_OFFLINE_TTS_KOKORO_AXCL_IMPL_H_ +#define SHERPA_ONNX_CSRC_AXCL_OFFLINE_TTS_KOKORO_AXCL_IMPL_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fst/extensions/far/far.h" +#include "kaldifst/csrc/kaldi-fst-io.h" +#include "kaldifst/csrc/text-normalizer.h" +#include "sherpa-onnx/csrc/file-utils.h" +#include "sherpa-onnx/csrc/kokoro-multi-lang-lexicon.h" +#include "sherpa-onnx/csrc/lexicon.h" +#include "sherpa-onnx/csrc/macros.h" +#include "sherpa-onnx/csrc/offline-tts-frontend.h" +#include "sherpa-onnx/csrc/offline-tts-impl.h" +#include "sherpa-onnx/csrc/offline-tts-kokoro-model.h" +#include "sherpa-onnx/csrc/piper-phonemize-lexicon.h" +#include "sherpa-onnx/csrc/text-utils.h" + +#ifdef SHERPA_ONNX_ENABLE_AXCL +#include "sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h" +#endif + +namespace sherpa_onnx { + +template +class OfflineTtsKokoroAxclImpl : public OfflineTtsImpl { + public: + explicit OfflineTtsKokoroAxclImpl(const OfflineTtsConfig &config) + : config_(config), model_(std::make_unique(config.model)) { + InitFrontend(); + + if (!config.rule_fsts.empty()) { + std::vector files; + SplitStringToVector(config.rule_fsts, ",", false, &files); + tn_list_.reserve(files.size()); + for (const auto &f : files) { + if (config.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("rule fst: %{public}s", f.c_str()); +#else + SHERPA_ONNX_LOGE("rule fst: %s", f.c_str()); +#endif + } + tn_list_.push_back(std::make_unique(f)); + } + } + + if (!config.rule_fars.empty()) { + if (config.model.debug) { + SHERPA_ONNX_LOGE("Loading FST archives"); + } + std::vector files; + SplitStringToVector(config.rule_fars, ",", false, &files); + + tn_list_.reserve(files.size() + tn_list_.size()); + + for (const auto &f : files) { + if (config.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("rule far: %{public}s", f.c_str()); +#else + SHERPA_ONNX_LOGE("rule far: %s", f.c_str()); +#endif + } + std::unique_ptr> reader( + fst::FarReader::Open(f)); + for (; !reader->Done(); reader->Next()) { + std::unique_ptr r( + fst::CastOrConvertToConstFst(reader->GetFst()->Copy())); + + tn_list_.push_back( + std::make_unique(std::move(r))); + } + } + + if (config.model.debug) { + SHERPA_ONNX_LOGE("FST archives loaded!"); + } + } + } + + template + OfflineTtsKokoroAxclImpl(Manager *mgr, const OfflineTtsConfig &config) + : config_(config), + model_(std::make_unique(mgr, config.model)) { + InitFrontend(mgr); + + if (!config.rule_fsts.empty()) { + std::vector files; + SplitStringToVector(config.rule_fsts, ",", false, &files); + tn_list_.reserve(files.size()); + for (const auto &f : files) { + if (config.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("rule fst: %{public}s", f.c_str()); +#else + SHERPA_ONNX_LOGE("rule fst: %s", f.c_str()); +#endif + } + auto buf = ReadFile(mgr, f); + std::istringstream is(std::string(buf.data(), buf.size())); + tn_list_.push_back(std::make_unique(is)); + } + } + + if (!config.rule_fars.empty()) { + std::vector files; + SplitStringToVector(config.rule_fars, ",", false, &files); + tn_list_.reserve(files.size() + tn_list_.size()); + + for (const auto &f : files) { + if (config.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("rule far: %{public}s", f.c_str()); +#else + SHERPA_ONNX_LOGE("rule far: %s", f.c_str()); +#endif + } + + auto buf = ReadFile(mgr, f); + + std::unique_ptr s( + new std::istringstream(std::string(buf.data(), buf.size()))); + + std::unique_ptr> reader( + fst::FarReader::Open(std::move(s))); + + for (; !reader->Done(); reader->Next()) { + std::unique_ptr r( + fst::CastOrConvertToConstFst(reader->GetFst()->Copy())); + + tn_list_.push_back( + std::make_unique(std::move(r))); + } // for (; !reader->Done(); reader->Next()) + } // for (const auto &f : files) + } // if (!config.rule_fars.empty()) + } + + int32_t SampleRate() const override { + return model_->GetMetaData().sample_rate; + } + + int32_t NumSpeakers() const override { + return model_->GetMetaData().num_speakers; + } + + GeneratedAudio Generate( + const std::string &_text, int64_t sid = 0, float speed = 1.0, + GeneratedAudioCallback callback = nullptr) const override { + const auto &meta_data = model_->GetMetaData(); + int32_t num_speakers = meta_data.num_speakers; + + if (num_speakers == 0 && sid != 0) { +#if __OHOS__ + SHERPA_ONNX_LOGE( + "This is a single-speaker model and supports only sid 0. Given sid: " + "%{public}d. sid is ignored", + static_cast(sid)); +#else + SHERPA_ONNX_LOGE( + "This is a single-speaker model and supports only sid 0. Given sid: " + "%d. sid is ignored", + static_cast(sid)); +#endif + } + + if (num_speakers != 0 && (sid >= num_speakers || sid < 0)) { +#if __OHOS__ + SHERPA_ONNX_LOGE( + "This model contains only %{public}d speakers. sid should be in the " + "range [%{public}d, %{public}d]. Given: %{public}d. Use sid=0", + num_speakers, 0, num_speakers - 1, static_cast(sid)); +#else + SHERPA_ONNX_LOGE( + "This model contains only %d speakers. sid should be in the range " + "[%d, %d]. Given: %d. Use sid=0", + num_speakers, 0, num_speakers - 1, static_cast(sid)); +#endif + sid = 0; + } + + std::string text = _text; + if (config_.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("Raw text: %{public}s", text.c_str()); +#else + SHERPA_ONNX_LOGE("Raw text: %s", text.c_str()); +#endif + std::ostringstream os; + os << "In bytes (hex):\n"; + const auto p = reinterpret_cast(text.c_str()); + for (int32_t i = 0; i != text.size(); ++i) { + os << std::setw(2) << std::setfill('0') << std::hex + << static_cast(p[i]) << " "; + } + os << "\n"; + +#if __OHOS__ + SHERPA_ONNX_LOGE("%{public}s", os.str().c_str()); +#else + SHERPA_ONNX_LOGE("%s", os.str().c_str()); +#endif + } + + if (!tn_list_.empty()) { + for (const auto &tn : tn_list_) { + text = tn->Normalize(text); + if (config_.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("After normalizing: %{public}s", text.c_str()); +#else + SHERPA_ONNX_LOGE("After normalizing: %s", text.c_str()); +#endif + } + } + } + + std::vector token_ids = frontend_->ConvertTextToTokenIds( + text, config_.model.kokoro.lang.empty() ? meta_data.voice + : config_.model.kokoro.lang); + + if (token_ids.empty() || + (token_ids.size() == 1 && token_ids[0].tokens.empty())) { +#if __OHOS__ + SHERPA_ONNX_LOGE("Failed to convert '%{public}s' to token IDs", + text.c_str()); +#else + SHERPA_ONNX_LOGE("Failed to convert '%s' to token IDs", text.c_str()); +#endif + return {}; + } + + std::vector> x; + + x.reserve(token_ids.size()); + + for (auto &i : token_ids) { + x.push_back(std::move(i.tokens)); + } + + int32_t x_size = static_cast(x.size()); + + if (config_.max_num_sentences != 1) { +#if __OHOS__ + SHERPA_ONNX_LOGE( + "max_num_sentences (%{public}d) != 1 is ignored for Kokoro TTS " + "models", + config_.max_num_sentences); +#else + SHERPA_ONNX_LOGE( + "max_num_sentences (%d) != 1 is ignored for Kokoro TTS models", + config_.max_num_sentences); +#endif + } + + // the input text is too long, we process sentences within it in batches + // to avoid OOM. Batch size is config_.max_num_sentences + std::vector> batch_x; + + int32_t batch_size = 1; + batch_x.reserve(config_.max_num_sentences); + int32_t num_batches = x_size / batch_size; + + if (config_.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE( + "Split it into %{public}d batches. batch size: " + "%{public}d. Number of sentences: %{public}d", + num_batches, batch_size, x_size); +#else + SHERPA_ONNX_LOGE( + "Split it into %d batches. batch size: %d. Number " + "of sentences: %d", + num_batches, batch_size, x_size); +#endif + } + + GeneratedAudio ans; + + int32_t should_continue = 1; + + int32_t k = 0; + + for (int32_t b = 0; b != num_batches && should_continue; ++b) { + batch_x.clear(); + for (int32_t i = 0; i != batch_size; ++i, ++k) { + batch_x.push_back(std::move(x[k])); + } + + auto audio = Process(batch_x, sid, speed); + ans.sample_rate = audio.sample_rate; + ans.samples.insert(ans.samples.end(), audio.samples.begin(), + audio.samples.end()); + if (callback) { + should_continue = callback(audio.samples.data(), audio.samples.size(), + (b + 1) * 1.0 / num_batches); + // Caution(fangjun): audio is freed when the callback returns, so users + // should copy the data if they want to access the data after + // the callback returns to avoid segmentation fault. + } + } + + batch_x.clear(); + while (k < static_cast(x.size()) && should_continue) { + batch_x.push_back(std::move(x[k])); + + ++k; + } + + if (!batch_x.empty()) { + auto audio = Process(batch_x, sid, speed); + ans.sample_rate = audio.sample_rate; + ans.samples.insert(ans.samples.end(), audio.samples.begin(), + audio.samples.end()); + if (callback) { + callback(audio.samples.data(), audio.samples.size(), 1.0); + // Caution(fangjun): audio is freed when the callback returns, so users + // should copy the data if they want to access the data after + // the callback returns to avoid segmentation fault. + } + } + + return ans; + } + + private: + template + void InitFrontend(Manager *mgr) { + const auto &meta_data = model_->GetMetaData(); + + if (meta_data.version >= 2) { + // this is a multi-lingual model, we require that you pass lexicon + if (config_.model.kokoro.lexicon.empty() && + config_.model.kokoro.lang.empty()) { + SHERPA_ONNX_LOGE("Current model version: '%d'", meta_data.version); + SHERPA_ONNX_LOGE( + "You are using a multi-lingual Kokoro model (e.g., Kokoro >= " + "v1.0). Please pass --kokoro-lexicon or provide --kokoro-lang"); + SHERPA_ONNX_EXIT(-1); + } + + frontend_ = std::make_unique( + mgr, config_.model.kokoro.tokens, config_.model.kokoro.lexicon, + config_.model.kokoro.data_dir, meta_data, config_.model.debug); + + return; + } + + frontend_ = std::make_unique( + mgr, config_.model.kokoro.tokens, config_.model.kokoro.data_dir, + meta_data); + } + + void InitFrontend() { + const auto &meta_data = model_->GetMetaData(); + if (meta_data.version >= 2) { + // this is a multi-lingual model, we require that you pass lexicon + if (config_.model.kokoro.lexicon.empty() && + config_.model.kokoro.lang.empty()) { + SHERPA_ONNX_LOGE("Current model version: '%d'", meta_data.version); + SHERPA_ONNX_LOGE( + "You are using a multi-lingual Kokoro model (e.g., Kokoro >= " + "v1.0). please pass --kokoro-lexicon or --kokoro-lang"); + SHERPA_ONNX_EXIT(-1); + } + + frontend_ = std::make_unique( + config_.model.kokoro.tokens, config_.model.kokoro.lexicon, + config_.model.kokoro.data_dir, meta_data, config_.model.debug); + + return; + } + + // this is for kokoro v0.19, which supports only English + frontend_ = std::make_unique( + config_.model.kokoro.tokens, config_.model.kokoro.data_dir, meta_data); + } + + GeneratedAudio Process(const std::vector> &tokens, + int32_t sid, float speed) const { + int32_t num_tokens = 0; + for (const auto &k : tokens) { + num_tokens += k.size(); + } + + std::vector x; + x.reserve(num_tokens); + for (const auto &k : tokens) { + x.insert(x.end(), k.begin(), k.end()); + } + + GeneratedAudio ans; + ans.sample_rate = model_->GetMetaData().sample_rate; + +#ifdef SHERPA_ONNX_ENABLE_AXCL + if constexpr (std::is_same_v) { + std::vector audio = model_->Run(x, sid, speed); + ans.samples = std::move(audio); + } else +#endif + { + auto memory_info = + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + + std::array x_shape = {1, static_cast(x.size())}; + Ort::Value x_tensor = Ort::Value::CreateTensor( + memory_info, x.data(), x.size(), x_shape.data(), x_shape.size()); + + Ort::Value audio = model_->Run(std::move(x_tensor), sid, speed); + + std::vector audio_shape = + audio.GetTensorTypeAndShapeInfo().GetShape(); + + int64_t total = 1; + // The output shape may be (1, 1, total) or (1, total) or (total,) + for (auto i : audio_shape) { + total *= i; + } + + const float *p = audio.GetTensorData(); + ans.samples = std::vector(p, p + total); + } + + float silence_scale = config_.silence_scale; + if (silence_scale != 1) { + ans = ans.ScaleSilence(silence_scale); + } + + return ans; + } + + private: + OfflineTtsConfig config_; + std::unique_ptr model_; + std::vector> tn_list_; + std::unique_ptr frontend_; +}; + +} // namespace sherpa_onnx +#endif // SHERPA_ONNX_CSRC_AXCL_OFFLINE_TTS_KOKORO_AXCL_IMPL_H_ diff --git a/sherpa-onnx/csrc/offline-tts-impl.cc b/sherpa-onnx/csrc/offline-tts-impl.cc index de1a0173da..84bffbd41e 100644 --- a/sherpa-onnx/csrc/offline-tts-impl.cc +++ b/sherpa-onnx/csrc/offline-tts-impl.cc @@ -25,7 +25,7 @@ #include "sherpa-onnx/csrc/offline-tts-zipvoice-impl.h" #ifdef SHERPA_ONNX_ENABLE_AXCL -#include "sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h" +#include "sherpa-onnx/csrc/axcl/offline-tts-kokoro-axcl-impl.h" #endif namespace sherpa_onnx { @@ -47,7 +47,7 @@ std::unique_ptr OfflineTtsImpl::Create( if (config.model.provider == "axcl") { #if SHERPA_ONNX_ENABLE_AXCL if (!config.model.kokoro.model.empty()) { - return std::make_unique>( + return std::make_unique>( config); } else { SHERPA_ONNX_LOGE( @@ -75,8 +75,7 @@ std::unique_ptr OfflineTtsImpl::Create( !config.model.zipvoice.decoder.empty()) { return std::make_unique(config); } else if (!config.model.kokoro.model.empty()) { - return std::make_unique>( - config); + return std::make_unique(config); } else if (!config.model.kitten.model.empty()) { return std::make_unique(config); } else if (!config.model.pocket.lm_flow.empty()) { @@ -101,8 +100,7 @@ std::unique_ptr OfflineTtsImpl::Create( !config.model.zipvoice.decoder.empty()) { return std::make_unique(mgr, config); } else if (!config.model.kokoro.model.empty()) { - return std::make_unique>( - mgr, config); + return std::make_unique(mgr, config); } else if (!config.model.kitten.model.empty()) { return std::make_unique(mgr, config); } else if (!config.model.pocket.lm_flow.empty()) { diff --git a/sherpa-onnx/csrc/offline-tts-kokoro-impl.h b/sherpa-onnx/csrc/offline-tts-kokoro-impl.h index 6febb001cb..0ca2921114 100644 --- a/sherpa-onnx/csrc/offline-tts-kokoro-impl.h +++ b/sherpa-onnx/csrc/offline-tts-kokoro-impl.h @@ -7,9 +7,8 @@ #include #include #include -#include #include -#include +#include #include #include @@ -26,17 +25,13 @@ #include "sherpa-onnx/csrc/piper-phonemize-lexicon.h" #include "sherpa-onnx/csrc/text-utils.h" -#ifdef SHERPA_ONNX_ENABLE_AXCL -#include "sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.h" -#endif - namespace sherpa_onnx { -template class OfflineTtsKokoroImpl : public OfflineTtsImpl { public: explicit OfflineTtsKokoroImpl(const OfflineTtsConfig &config) - : config_(config), model_(std::make_unique(config.model)) { + : config_(config), + model_(std::make_unique(config.model)) { InitFrontend(); if (!config.rule_fsts.empty()) { @@ -92,7 +87,7 @@ class OfflineTtsKokoroImpl : public OfflineTtsImpl { template OfflineTtsKokoroImpl(Manager *mgr, const OfflineTtsConfig &config) : config_(config), - model_(std::make_unique(mgr, config.model)) { + model_(std::make_unique(mgr, config.model)) { InitFrontend(mgr); if (!config.rule_fsts.empty()) { @@ -142,8 +137,8 @@ class OfflineTtsKokoroImpl : public OfflineTtsImpl { tn_list_.push_back( std::make_unique(std::move(r))); } // for (; !reader->Done(); reader->Next()) - } // for (const auto &f : files) - } // if (!config.rule_fars.empty()) + } // for (const auto &f : files) + } // if (!config.rule_fars.empty()) } int32_t SampleRate() const override { @@ -399,38 +394,30 @@ class OfflineTtsKokoroImpl : public OfflineTtsImpl { x.insert(x.end(), k.begin(), k.end()); } - GeneratedAudio ans; - ans.sample_rate = model_->GetMetaData().sample_rate; + auto memory_info = + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); -#ifdef SHERPA_ONNX_ENABLE_AXCL - if constexpr (std::is_same_v) { - std::vector audio = model_->Run(x, sid, speed); - ans.samples = std::move(audio); - } else -#endif - { - auto memory_info = - Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); - - std::array x_shape = {1, static_cast(x.size())}; - Ort::Value x_tensor = Ort::Value::CreateTensor( - memory_info, x.data(), x.size(), x_shape.data(), x_shape.size()); + std::array x_shape = {1, static_cast(x.size())}; + Ort::Value x_tensor = Ort::Value::CreateTensor( + memory_info, x.data(), x.size(), x_shape.data(), x_shape.size()); Ort::Value audio = model_->Run(std::move(x_tensor), sid, speed); - std::vector audio_shape = - audio.GetTensorTypeAndShapeInfo().GetShape(); - - int64_t total = 1; - // The output shape may be (1, 1, total) or (1, total) or (total,) - for (auto i : audio_shape) { - total *= i; - } + std::vector audio_shape = + audio.GetTensorTypeAndShapeInfo().GetShape(); - const float *p = audio.GetTensorData(); - ans.samples = std::vector(p, p + total); + int64_t total = 1; + // The output shape may be (1, 1, total) or (1, total) or (total,) + for (auto i : audio_shape) { + total *= i; } + const float *p = audio.GetTensorData(); + + GeneratedAudio ans; + ans.sample_rate = model_->GetMetaData().sample_rate; + ans.samples = std::vector(p, p + total); + float silence_scale = config_.silence_scale; if (silence_scale != 1) { ans = ans.ScaleSilence(silence_scale); @@ -441,7 +428,7 @@ class OfflineTtsKokoroImpl : public OfflineTtsImpl { private: OfflineTtsConfig config_; - std::unique_ptr model_; + std::unique_ptr model_; std::vector> tn_list_; std::unique_ptr frontend_; }; From 1b940fadc1fd5f17064b693c97112da841ed9a91 Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Wed, 11 Mar 2026 11:22:19 +0800 Subject: [PATCH 03/10] kokoro TTS Supported axera backend --- sherpa-onnx/csrc/CMakeLists.txt | 5 + .../axera/offline-tts-kokoro-axera-impl.h | 450 +++++++++++ .../axera/offline-tts-kokoro-model-axera.cc | 747 ++++++++++++++++++ .../axera/offline-tts-kokoro-model-axera.h | 38 + sherpa-onnx/csrc/offline-tts-impl.cc | 75 +- 5 files changed, 1313 insertions(+), 2 deletions(-) create mode 100644 sherpa-onnx/csrc/axera/offline-tts-kokoro-axera-impl.h create mode 100644 sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc create mode 100644 sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.h diff --git a/sherpa-onnx/csrc/CMakeLists.txt b/sherpa-onnx/csrc/CMakeLists.txt index 52b106c942..db9a8c3917 100755 --- a/sherpa-onnx/csrc/CMakeLists.txt +++ b/sherpa-onnx/csrc/CMakeLists.txt @@ -211,6 +211,11 @@ if(SHERPA_ONNX_ENABLE_AXERA) ./axera/offline-sense-voice-model-axera.cc ./axera/utils.cc ) + if(SHERPA_ONNX_ENABLE_TTS) + list(APPEND sources + ./axera/offline-tts-kokoro-model-axera.cc + ) + endif() endif() if(SHERPA_ONNX_ENABLE_AXCL) diff --git a/sherpa-onnx/csrc/axera/offline-tts-kokoro-axera-impl.h b/sherpa-onnx/csrc/axera/offline-tts-kokoro-axera-impl.h new file mode 100644 index 0000000000..7bf533db66 --- /dev/null +++ b/sherpa-onnx/csrc/axera/offline-tts-kokoro-axera-impl.h @@ -0,0 +1,450 @@ +// sherpa-onnx/csrc/axera/offline-tts-kokoro-axera-impl.h +// +// Copyright (c) 2025 Xiaomi Corporation +#ifndef SHERPA_ONNX_CSRC_AXERA_OFFLINE_TTS_KOKORO_AXERA_IMPL_H_ +#define SHERPA_ONNX_CSRC_AXERA_OFFLINE_TTS_KOKORO_AXERA_IMPL_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fst/extensions/far/far.h" +#include "kaldifst/csrc/kaldi-fst-io.h" +#include "kaldifst/csrc/text-normalizer.h" +#include "sherpa-onnx/csrc/file-utils.h" +#include "sherpa-onnx/csrc/kokoro-multi-lang-lexicon.h" +#include "sherpa-onnx/csrc/lexicon.h" +#include "sherpa-onnx/csrc/macros.h" +#include "sherpa-onnx/csrc/offline-tts-frontend.h" +#include "sherpa-onnx/csrc/offline-tts-impl.h" +#include "sherpa-onnx/csrc/offline-tts-kokoro-model.h" +#include "sherpa-onnx/csrc/piper-phonemize-lexicon.h" +#include "sherpa-onnx/csrc/text-utils.h" + +#ifdef SHERPA_ONNX_ENABLE_AXERA +#include "sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.h" +#endif + +namespace sherpa_onnx { + +template +class OfflineTtsKokoroAxeraImpl : public OfflineTtsImpl { + public: + explicit OfflineTtsKokoroAxeraImpl(const OfflineTtsConfig &config) + : config_(config), model_(std::make_unique(config.model)) { + InitFrontend(); + + if (!config.rule_fsts.empty()) { + std::vector files; + SplitStringToVector(config.rule_fsts, ",", false, &files); + tn_list_.reserve(files.size()); + for (const auto &f : files) { + if (config.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("rule fst: %{public}s", f.c_str()); +#else + SHERPA_ONNX_LOGE("rule fst: %s", f.c_str()); +#endif + } + tn_list_.push_back(std::make_unique(f)); + } + } + + if (!config.rule_fars.empty()) { + if (config.model.debug) { + SHERPA_ONNX_LOGE("Loading FST archives"); + } + std::vector files; + SplitStringToVector(config.rule_fars, ",", false, &files); + + tn_list_.reserve(files.size() + tn_list_.size()); + + for (const auto &f : files) { + if (config.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("rule far: %{public}s", f.c_str()); +#else + SHERPA_ONNX_LOGE("rule far: %s", f.c_str()); +#endif + } + std::unique_ptr> reader( + fst::FarReader::Open(f)); + for (; !reader->Done(); reader->Next()) { + std::unique_ptr r( + fst::CastOrConvertToConstFst(reader->GetFst()->Copy())); + + tn_list_.push_back( + std::make_unique(std::move(r))); + } + } + + if (config.model.debug) { + SHERPA_ONNX_LOGE("FST archives loaded!"); + } + } + } + + template + OfflineTtsKokoroAxeraImpl(Manager *mgr, const OfflineTtsConfig &config) + : config_(config), + model_(std::make_unique(mgr, config.model)) { + InitFrontend(mgr); + + if (!config.rule_fsts.empty()) { + std::vector files; + SplitStringToVector(config.rule_fsts, ",", false, &files); + tn_list_.reserve(files.size()); + for (const auto &f : files) { + if (config.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("rule fst: %{public}s", f.c_str()); +#else + SHERPA_ONNX_LOGE("rule fst: %s", f.c_str()); +#endif + } + auto buf = ReadFile(mgr, f); + std::istringstream is(std::string(buf.data(), buf.size())); + tn_list_.push_back(std::make_unique(is)); + } + } + + if (!config.rule_fars.empty()) { + std::vector files; + SplitStringToVector(config.rule_fars, ",", false, &files); + tn_list_.reserve(files.size() + tn_list_.size()); + + for (const auto &f : files) { + if (config.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("rule far: %{public}s", f.c_str()); +#else + SHERPA_ONNX_LOGE("rule far: %s", f.c_str()); +#endif + } + + auto buf = ReadFile(mgr, f); + + std::unique_ptr s( + new std::istringstream(std::string(buf.data(), buf.size()))); + + std::unique_ptr> reader( + fst::FarReader::Open(std::move(s))); + + for (; !reader->Done(); reader->Next()) { + std::unique_ptr r( + fst::CastOrConvertToConstFst(reader->GetFst()->Copy())); + + tn_list_.push_back( + std::make_unique(std::move(r))); + } // for (; !reader->Done(); reader->Next()) + } // for (const auto &f : files) + } // if (!config.rule_fars.empty()) + } + + int32_t SampleRate() const override { + return model_->GetMetaData().sample_rate; + } + + int32_t NumSpeakers() const override { + return model_->GetMetaData().num_speakers; + } + + GeneratedAudio Generate( + const std::string &_text, int64_t sid = 0, float speed = 1.0, + GeneratedAudioCallback callback = nullptr) const override { + const auto &meta_data = model_->GetMetaData(); + int32_t num_speakers = meta_data.num_speakers; + + if (num_speakers == 0 && sid != 0) { +#if __OHOS__ + SHERPA_ONNX_LOGE( + "This is a single-speaker model and supports only sid 0. Given sid: " + "%{public}d. sid is ignored", + static_cast(sid)); +#else + SHERPA_ONNX_LOGE( + "This is a single-speaker model and supports only sid 0. Given sid: " + "%d. sid is ignored", + static_cast(sid)); +#endif + } + + if (num_speakers != 0 && (sid >= num_speakers || sid < 0)) { +#if __OHOS__ + SHERPA_ONNX_LOGE( + "This model contains only %{public}d speakers. sid should be in the " + "range [%{public}d, %{public}d]. Given: %{public}d. Use sid=0", + num_speakers, 0, num_speakers - 1, static_cast(sid)); +#else + SHERPA_ONNX_LOGE( + "This model contains only %d speakers. sid should be in the range " + "[%d, %d]. Given: %d. Use sid=0", + num_speakers, 0, num_speakers - 1, static_cast(sid)); +#endif + sid = 0; + } + + std::string text = _text; + if (config_.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("Raw text: %{public}s", text.c_str()); +#else + SHERPA_ONNX_LOGE("Raw text: %s", text.c_str()); +#endif + std::ostringstream os; + os << "In bytes (hex):\n"; + const auto p = reinterpret_cast(text.c_str()); + for (int32_t i = 0; i != text.size(); ++i) { + os << std::setw(2) << std::setfill('0') << std::hex + << static_cast(p[i]) << " "; + } + os << "\n"; + +#if __OHOS__ + SHERPA_ONNX_LOGE("%{public}s", os.str().c_str()); +#else + SHERPA_ONNX_LOGE("%s", os.str().c_str()); +#endif + } + + if (!tn_list_.empty()) { + for (const auto &tn : tn_list_) { + text = tn->Normalize(text); + if (config_.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE("After normalizing: %{public}s", text.c_str()); +#else + SHERPA_ONNX_LOGE("After normalizing: %s", text.c_str()); +#endif + } + } + } + + std::vector token_ids = frontend_->ConvertTextToTokenIds( + text, config_.model.kokoro.lang.empty() ? meta_data.voice + : config_.model.kokoro.lang); + + if (token_ids.empty() || + (token_ids.size() == 1 && token_ids[0].tokens.empty())) { +#if __OHOS__ + SHERPA_ONNX_LOGE("Failed to convert '%{public}s' to token IDs", + text.c_str()); +#else + SHERPA_ONNX_LOGE("Failed to convert '%s' to token IDs", text.c_str()); +#endif + return {}; + } + + std::vector> x; + + x.reserve(token_ids.size()); + + for (auto &i : token_ids) { + x.push_back(std::move(i.tokens)); + } + + int32_t x_size = static_cast(x.size()); + + if (config_.max_num_sentences != 1) { +#if __OHOS__ + SHERPA_ONNX_LOGE( + "max_num_sentences (%{public}d) != 1 is ignored for Kokoro TTS " + "models", + config_.max_num_sentences); +#else + SHERPA_ONNX_LOGE( + "max_num_sentences (%d) != 1 is ignored for Kokoro TTS models", + config_.max_num_sentences); +#endif + } + + // the input text is too long, we process sentences within it in batches + // to avoid OOM. Batch size is config_.max_num_sentences + std::vector> batch_x; + + int32_t batch_size = 1; + batch_x.reserve(config_.max_num_sentences); + int32_t num_batches = x_size / batch_size; + + if (config_.model.debug) { +#if __OHOS__ + SHERPA_ONNX_LOGE( + "Split it into %{public}d batches. batch size: " + "%{public}d. Number of sentences: %{public}d", + num_batches, batch_size, x_size); +#else + SHERPA_ONNX_LOGE( + "Split it into %d batches. batch size: %d. Number " + "of sentences: %d", + num_batches, batch_size, x_size); +#endif + } + + GeneratedAudio ans; + + int32_t should_continue = 1; + + int32_t k = 0; + + for (int32_t b = 0; b != num_batches && should_continue; ++b) { + batch_x.clear(); + for (int32_t i = 0; i != batch_size; ++i, ++k) { + batch_x.push_back(std::move(x[k])); + } + + auto audio = Process(batch_x, sid, speed); + ans.sample_rate = audio.sample_rate; + ans.samples.insert(ans.samples.end(), audio.samples.begin(), + audio.samples.end()); + if (callback) { + should_continue = callback(audio.samples.data(), audio.samples.size(), + (b + 1) * 1.0 / num_batches); + // Caution(fangjun): audio is freed when the callback returns, so users + // should copy the data if they want to access the data after + // the callback returns to avoid segmentation fault. + } + } + + batch_x.clear(); + while (k < static_cast(x.size()) && should_continue) { + batch_x.push_back(std::move(x[k])); + + ++k; + } + + if (!batch_x.empty()) { + auto audio = Process(batch_x, sid, speed); + ans.sample_rate = audio.sample_rate; + ans.samples.insert(ans.samples.end(), audio.samples.begin(), + audio.samples.end()); + if (callback) { + callback(audio.samples.data(), audio.samples.size(), 1.0); + // Caution(fangjun): audio is freed when the callback returns, so users + // should copy the data if they want to access the data after + // the callback returns to avoid segmentation fault. + } + } + + return ans; + } + + private: + template + void InitFrontend(Manager *mgr) { + const auto &meta_data = model_->GetMetaData(); + + if (meta_data.version >= 2) { + // this is a multi-lingual model, we require that you pass lexicon + if (config_.model.kokoro.lexicon.empty() && + config_.model.kokoro.lang.empty()) { + SHERPA_ONNX_LOGE("Current model version: '%d'", meta_data.version); + SHERPA_ONNX_LOGE( + "You are using a multi-lingual Kokoro model (e.g., Kokoro >= " + "v1.0). Please pass --kokoro-lexicon or provide --kokoro-lang"); + SHERPA_ONNX_EXIT(-1); + } + + frontend_ = std::make_unique( + mgr, config_.model.kokoro.tokens, config_.model.kokoro.lexicon, + config_.model.kokoro.data_dir, meta_data, config_.model.debug); + + return; + } + + frontend_ = std::make_unique( + mgr, config_.model.kokoro.tokens, config_.model.kokoro.data_dir, + meta_data); + } + + void InitFrontend() { + const auto &meta_data = model_->GetMetaData(); + if (meta_data.version >= 2) { + // this is a multi-lingual model, we require that you pass lexicon + if (config_.model.kokoro.lexicon.empty() && + config_.model.kokoro.lang.empty()) { + SHERPA_ONNX_LOGE("Current model version: '%d'", meta_data.version); + SHERPA_ONNX_LOGE( + "You are using a multi-lingual Kokoro model (e.g., Kokoro >= " + "v1.0). please pass --kokoro-lexicon or --kokoro-lang"); + SHERPA_ONNX_EXIT(-1); + } + + frontend_ = std::make_unique( + config_.model.kokoro.tokens, config_.model.kokoro.lexicon, + config_.model.kokoro.data_dir, meta_data, config_.model.debug); + + return; + } + + // this is for kokoro v0.19, which supports only English + frontend_ = std::make_unique( + config_.model.kokoro.tokens, config_.model.kokoro.data_dir, meta_data); + } + + GeneratedAudio Process(const std::vector> &tokens, + int32_t sid, float speed) const { + int32_t num_tokens = 0; + for (const auto &k : tokens) { + num_tokens += k.size(); + } + + std::vector x; + x.reserve(num_tokens); + for (const auto &k : tokens) { + x.insert(x.end(), k.begin(), k.end()); + } + + GeneratedAudio ans; + ans.sample_rate = model_->GetMetaData().sample_rate; + +#ifdef SHERPA_ONNX_ENABLE_AXERA + if constexpr (std::is_same_v) { + std::vector audio = model_->Run(x, sid, speed); + ans.samples = std::move(audio); + } else +#endif + { + auto memory_info = + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + + std::array x_shape = {1, static_cast(x.size())}; + Ort::Value x_tensor = Ort::Value::CreateTensor( + memory_info, x.data(), x.size(), x_shape.data(), x_shape.size()); + + Ort::Value audio = model_->Run(std::move(x_tensor), sid, speed); + + std::vector audio_shape = + audio.GetTensorTypeAndShapeInfo().GetShape(); + + int64_t total = 1; + // The output shape may be (1, 1, total) or (1, total) or (total,) + for (auto i : audio_shape) { + total *= i; + } + + const float *p = audio.GetTensorData(); + ans.samples = std::vector(p, p + total); + } + + float silence_scale = config_.silence_scale; + if (silence_scale != 1) { + ans = ans.ScaleSilence(silence_scale); + } + + return ans; + } + + private: + OfflineTtsConfig config_; + std::unique_ptr model_; + std::vector> tn_list_; + std::unique_ptr frontend_; +}; + +} // namespace sherpa_onnx +#endif // SHERPA_ONNX_CSRC_AXERA_OFFLINE_TTS_KOKORO_AXERA_IMPL_H_ \ No newline at end of file diff --git a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc new file mode 100644 index 0000000000..7699f1aedd --- /dev/null +++ b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc @@ -0,0 +1,747 @@ +// sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc +// +// Copyright (c) 2025 Xiaomi Corporation + +#include "sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if __ANDROID_API__ >= 9 +#include "android/asset_manager.h" +#include "android/asset_manager_jni.h" +#endif + +#if __OHOS__ +#include "rawfile/raw_file_manager.h" +#endif + +#include "ax_engine_api.h" // NOLINT +#include "ax_sys_api.h" // NOLINT +#include "kaldi-native-fbank/csrc/istft.h" +#include "sherpa-onnx/csrc/axera/ax-engine-guard.h" +#include "sherpa-onnx/csrc/axera/utils.h" +#include "sherpa-onnx/csrc/file-utils.h" +#include "sherpa-onnx/csrc/macros.h" +#include "sherpa-onnx/csrc/onnx-utils.h" +#include "sherpa-onnx/csrc/session.h" + +namespace sherpa_onnx { + +static std::vector Sigmoid(const std::vector &x) { + std::vector result(x.size()); + for (size_t i = 0; i < x.size(); i++) { + result[i] = 1.0f / (1.0f + std::exp(-x[i])); + } + return result; +} + +template +static std::vector Argsort(const std::vector &v, int32_t len, + bool reverse) { + std::vector idx(len); + std::iota(idx.begin(), idx.end(), 0); + + if (!reverse) { + std::stable_sort(idx.begin(), idx.end(), + [&v](size_t i1, size_t i2) { return v[i1] < v[i2]; }); + } else { + std::stable_sort(idx.begin(), idx.end(), + [&v](size_t i1, size_t i2) { return v[i1] > v[i2]; }); + } + + return idx; +} + +template +static std::vector NpRepeat(const std::vector &v, + const std::vector ×) { + std::vector result; + for (size_t i = 0; i < times.size(); i++) { + for (int32_t n = 0; n < times[i]; n++) { + result.push_back(v[i]); + } + } + return result; +} + +class AxeraModel { + public: + explicit AxeraModel(const std::string &filename) { + auto buf = ReadFile(filename); + Init(buf.data(), buf.size()); + } + + AxeraModel(const void *cpu_buf, size_t buf_len_in_bytes) { + Init(cpu_buf, buf_len_in_bytes); + } + + ~AxeraModel() { + if (io_data_.pInputs || io_data_.pOutputs) { + FreeIO(&io_data_); + } + + if (handle_) { + AX_ENGINE_DestroyHandle(handle_); + } + } + + std::vector TensorShape(const std::string &name) const { + auto it = input_name_to_index_.find(name); + if (it != input_name_to_index_.end()) { + return input_tensor_shapes_[it->second]; + } + + auto it2 = output_name_to_index_.find(name); + if (it2 != output_name_to_index_.end()) { + return output_tensor_shapes_[it2->second]; + } + + SHERPA_ONNX_LOGE("Found no tensor with name: '%s'", name.c_str()); + return {}; + } + + template + bool SetInputTensorData(const std::string &name, const T *p, int32_t n) { + auto it = input_name_to_index_.find(name); + if (it == input_name_to_index_.end()) { + SHERPA_ONNX_LOGE("Found no input tensor with name: '%s'", name.c_str()); + return false; + } + + int32_t i = it->second; + size_t expected_size = io_info_->pInputs[i].nSize; + size_t given_size = n * sizeof(T); + + if (expected_size != given_size) { + SHERPA_ONNX_LOGE( + "Input tensor '%s' size mismatch. Expected %zu bytes, got %zu bytes", + name.c_str(), expected_size, given_size); + return false; + } + + std::memcpy(io_data_.pInputs[i].pVirAddr, p, expected_size); + return true; + } + + std::vector GetOutputTensorData(const std::string &name) { + auto it = output_name_to_index_.find(name); + if (it == output_name_to_index_.end()) { + SHERPA_ONNX_LOGE("Found no output tensor with name: '%s'", name.c_str()); + return {}; + } + + int32_t i = it->second; + const auto &out_meta = io_info_->pOutputs[i]; + auto &out_buf = io_data_.pOutputs[i]; + + AX_SYS_MinvalidateCache(out_buf.phyAddr, out_buf.pVirAddr, out_meta.nSize); + + size_t out_elems = out_meta.nSize / sizeof(float); + std::vector out(out_elems); + std::memcpy(out.data(), out_buf.pVirAddr, out_meta.nSize); + + return out; + } + + bool Run() { + auto ret = AX_ENGINE_RunSync(handle_, &io_data_); + if (ret != 0) { + SHERPA_ONNX_LOGE("AX_ENGINE_RunSync failed, ret = %d", ret); + return false; + } + + return true; + } + + private: + void Init(const void *model_data, size_t model_data_length) { + InitContext(model_data, model_data_length, false, &handle_); + InitInputOutputAttrs(handle_, false, &io_info_); + PrepareIO(io_info_, &io_data_, false); + + input_tensor_shapes_.reserve(io_info_->nInputSize); + for (uint32_t i = 0; i < io_info_->nInputSize; ++i) { + const auto &in = io_info_->pInputs[i]; + std::string name = in.pName; + input_name_to_index_[name] = i; + input_tensor_shapes_.emplace_back(in.pShape, in.pShape + in.nShapeSize); + } + + output_tensor_shapes_.reserve(io_info_->nOutputSize); + for (uint32_t i = 0; i < io_info_->nOutputSize; ++i) { + const auto &out = io_info_->pOutputs[i]; + std::string name = out.pName; + output_name_to_index_[name] = i; + output_tensor_shapes_.emplace_back(out.pShape, + out.pShape + out.nShapeSize); + } + } + + private: + AX_ENGINE_HANDLE handle_ = nullptr; + AX_ENGINE_IO_INFO_T *io_info_ = nullptr; + AX_ENGINE_IO_T io_data_{}; + + std::unordered_map input_name_to_index_; + std::unordered_map output_name_to_index_; + + std::vector> input_tensor_shapes_; + std::vector> output_tensor_shapes_; +}; + +static std::string GetModelPath(const std::string &model, const std::string &p, + bool is_onnx = false) { + std::string suffix = is_onnx ? ".onnx" : ".axmodel"; + + std::string path1 = model + "/kokoro_" + p + suffix; + if (FileExists(path1)) { + return path1; + } + + return model + "_" + p + suffix; +} + +class OfflineTtsKokoroModelAxera::Impl { + private: + int32_t max_seq_len_ = 0; + std::vector styles_; + + public: + explicit Impl(const OfflineTtsModelConfig &config) + : config_(config), + env_(ORT_LOGGING_LEVEL_ERROR), + sess_opts_(GetSessionOptions(config)), + allocator_{} { + Init(config); + } + + template + Impl(Manager *mgr, const OfflineTtsModelConfig &config) + : config_(config), + env_(ORT_LOGGING_LEVEL_ERROR), + sess_opts_(GetSessionOptions(config)), + allocator_{} { + Init(mgr, config); + } + + const OfflineTtsKokoroModelMetaData &GetMetaData() const { + return meta_data_; + } + + std::vector Run(std::vector input_ids, int64_t sid, + float speed) { + std::lock_guard lock(mutex_); + + if (input_ids.size() < 2) { + return {}; + } + + if (config_.kokoro.length_scale != 1 && speed == 1) { + speed = 1. / config_.kokoro.length_scale; + } + + int32_t phoneme_len = static_cast(input_ids.size()) - 2; + if (phoneme_len < 0) phoneme_len = 0; + + std::vector ref_s = + LoadVoiceEmbedding(static_cast(sid), phoneme_len); + + std::vector audio; + bool success = RunModels(input_ids, ref_s, speed, audio); + if (!success) { + SHERPA_ONNX_LOGE("Run models failed!"); + return {}; + } + + return audio; + } + + private: + void Init(const OfflineTtsModelConfig &config) { + std::string model1_path = GetModelPath(config.kokoro.model, "part1"); + std::string model2_path = GetModelPath(config.kokoro.model, "part2"); + std::string model3_path = GetModelPath(config.kokoro.model, "part3"); + std::string model4_path = GetModelPath(config.kokoro.model, "part4", true); + + model1_ = std::make_unique(model1_path); + model2_ = std::make_unique(model2_path); + model3_ = std::make_unique(model3_path); + + auto buf4 = ReadFile(model4_path); + model4_ = std::make_unique(env_, buf4.data(), buf4.size(), + sess_opts_); + GetInputNames(model4_.get(), &model4_input_names_, + &model4_input_names_ptr_); + GetOutputNames(model4_.get(), &model4_output_names_, + &model4_output_names_ptr_); + + auto voices_buf = ReadFile(config.kokoro.voices); + LoadVoices(voices_buf.data(), voices_buf.size()); + } + + template + void Init(Manager *mgr, const OfflineTtsModelConfig &config) { + std::string model1_path = GetModelPath(config.kokoro.model, "part1"); + std::string model2_path = GetModelPath(config.kokoro.model, "part2"); + std::string model3_path = GetModelPath(config.kokoro.model, "part3"); + std::string model4_path = GetModelPath(config.kokoro.model, "part4", true); + + auto buf1 = ReadFile(mgr, model1_path); + model1_ = std::make_unique(buf1.data(), buf1.size()); + + auto buf2 = ReadFile(mgr, model2_path); + model2_ = std::make_unique(buf2.data(), buf2.size()); + + auto buf3 = ReadFile(mgr, model3_path); + model3_ = std::make_unique(buf3.data(), buf3.size()); + + auto buf4 = ReadFile(mgr, model4_path); + model4_ = std::make_unique(env_, buf4.data(), buf4.size(), + sess_opts_); + GetInputNames(model4_.get(), &model4_input_names_, + &model4_input_names_ptr_); + GetOutputNames(model4_.get(), &model4_output_names_, + &model4_output_names_ptr_); + + auto voices_buf = ReadFile(mgr, config.kokoro.voices); + LoadVoices(voices_buf.data(), voices_buf.size()); + } + + void LoadVoices(const char *voices_data, size_t voices_data_length) { + // get metadata + meta_data_.sample_rate = 24000; + meta_data_.num_speakers = voices_data_length / (sizeof(float) * 256 * 510); + meta_data_.version = 1; + meta_data_.has_espeak = 1; + meta_data_.voice = "en-us"; + meta_data_.max_token_len = 510; + + int32_t actual_num_floats = voices_data_length / sizeof(float); + styles_ = std::vector( + reinterpret_cast(voices_data), + reinterpret_cast(voices_data) + actual_num_floats); + + auto shape1 = model1_->TensorShape("input_ids"); + if (shape1.size() < 2) { + SHERPA_ONNX_LOGE("Unexpected shape rank for input_ids: %d", + static_cast(shape1.size())); + SHERPA_ONNX_EXIT(-1); + } + + max_seq_len_ = shape1[1]; + } + + std::vector LoadVoiceEmbedding(int32_t sid, int32_t phoneme_len) { + int32_t style_dim0 = 510; + int32_t style_dim1 = 256; + phoneme_len = std::max(phoneme_len, 0); + + sid = std::max(sid, 0); + if (meta_data_.num_speakers > 0) { + sid = std::min(sid, meta_data_.num_speakers - 1); + } + + std::vector ref_s(style_dim1); + if (phoneme_len < style_dim0) { + const float *p = styles_.data() + sid * style_dim0 * style_dim1 + + phoneme_len * style_dim1; + std::copy(p, p + style_dim1, ref_s.begin()); + } else { + int32_t idx = style_dim0 / 2; + const float *p = + styles_.data() + sid * style_dim0 * style_dim1 + idx * style_dim1; + std::copy(p, p + style_dim1, ref_s.begin()); + } + return ref_s; + } + + bool RunModels(std::vector &input_ids, + const std::vector &ref_s, float speed, + std::vector &audio) { + int32_t actual_len = input_ids.size(); + if (actual_len > max_seq_len_) { + SHERPA_ONNX_LOGE( + "Input ids length %d exceeds max_seq_len %d, truncating.", actual_len, + max_seq_len_); + input_ids.resize(max_seq_len_); + actual_len = max_seq_len_; + } + + int32_t padding_len = max_seq_len_ - actual_len; + if (padding_len > 0) { + std::vector padding(padding_len, 0); + input_ids.insert(input_ids.end(), padding.begin(), padding.end()); + } + + int32_t actual_content_frames = 0; + int32_t total_frames = 0; + if (!InferenceSingleChunk(input_ids, ref_s, actual_len, speed, audio, + actual_content_frames, total_frames)) { + return false; + } + + TrimAudioByContent(audio, actual_content_frames, total_frames, actual_len); + + return true; + } + + bool InferenceSingleChunk(std::vector &input_ids, + const std::vector &ref_s, int32_t actual_len, + float speed, std::vector &audio, + int32_t &actual_content_frames, + int32_t &total_frames) { + bool is_doubled = false; + int32_t original_actual_len = actual_len; + + PrepareInputIds(input_ids, actual_len, is_doubled); + + std::vector input_ids_i32(input_ids.begin(), input_ids.end()); + + std::vector text_mask; + ComputeExternalPreprocessing(actual_len, text_mask); + + std::vector text_mask_u8(text_mask.begin(), text_mask.end()); + + if (!model1_->SetInputTensorData("input_ids", input_ids_i32.data(), + input_ids_i32.size()) || + !model1_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()) || + !model1_->SetInputTensorData("text_mask", text_mask_u8.data(), + text_mask_u8.size())) { + return false; + } + + if (!model1_->Run()) { + return false; + } + + std::vector duration = model1_->GetOutputTensorData("duration"); + std::vector d = model1_->GetOutputTensorData("d"); + + std::vector pred_dur; + ProcessDuration(duration, actual_len, speed, pred_dur, total_frames); + + std::vector pred_aln_trg = + CreateAlignmentMatrix(pred_dur, total_frames); + + int32_t d_shape_2 = 640; + std::vector en(total_frames * d_shape_2, 0.0f); + + for (int32_t j = 0; j < d_shape_2; ++j) { + for (int32_t k = 0; k < total_frames; ++k) { + float sum = 0.0f; + for (int32_t i = 0; i < max_seq_len_; ++i) { + sum += d[i * d_shape_2 + j] * pred_aln_trg[i * total_frames + k]; + } + en[j * total_frames + k] = sum; + } + } + + std::vector text_mask_float(text_mask.begin(), text_mask.end()); + + if (!model2_->SetInputTensorData("en", en.data(), en.size()) || + !model2_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()) || + !model2_->SetInputTensorData("input_ids", input_ids_i32.data(), + input_ids_i32.size()) || + !model2_->SetInputTensorData("text_mask", text_mask_float.data(), + text_mask_float.size()) || + !model2_->SetInputTensorData("pred_aln_trg", pred_aln_trg.data(), + pred_aln_trg.size())) { + return false; + } + + if (!model2_->Run()) { + return false; + } + + std::vector f0_pred = model2_->GetOutputTensorData("F0_pred"); + std::vector n_pred = model2_->GetOutputTensorData("N_pred"); + std::vector asr = model2_->GetOutputTensorData("asr"); + + auto memory_info = + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + std::array f0_shape = {1, static_cast(f0_pred.size())}; + Ort::Value f0_tensor = + Ort::Value::CreateTensor(memory_info, f0_pred.data(), f0_pred.size(), + f0_shape.data(), f0_shape.size()); + + auto out4 = model4_->Run({}, model4_input_names_ptr_.data(), &f0_tensor, 1, + model4_output_names_ptr_.data(), + model4_output_names_ptr_.size()); + float *p_har = out4[0].GetTensorMutableData(); + auto har_shape = out4[0].GetTensorTypeAndShapeInfo().GetShape(); + int64_t har_num_elements = 1; + for (auto d : har_shape) { + har_num_elements *= d; + } + std::vector har(p_har, p_har + har_num_elements); + + if (!model3_->SetInputTensorData("asr", asr.data(), asr.size()) || + !model3_->SetInputTensorData("F0_pred", f0_pred.data(), + f0_pred.size()) || + !model3_->SetInputTensorData("N_pred", n_pred.data(), n_pred.size()) || + !model3_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()) || + !model3_->SetInputTensorData("har", har.data(), har.size())) { + return false; + } + + if (!model3_->Run()) { + return false; + } + + std::vector x_out = model3_->GetOutputTensorData("x"); + PostprocessXToAudio(x_out, 23041, audio); + + if (is_doubled) { + actual_content_frames = std::accumulate( + pred_dur.begin(), pred_dur.begin() + original_actual_len, 0); + audio.erase(audio.begin() + audio.size() / 2, audio.end()); + total_frames = total_frames / 2; + } else { + actual_content_frames = + std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); + } + return true; + } + + void ProcessDuration(const std::vector &duration, int32_t actual_len, + float speed, std::vector &pred_dur, + int32_t &total_frames) { + std::vector pred_dur_original(actual_len, 0); + std::vector duration_processed = Sigmoid(duration); + + int32_t dur_shape_2 = 50; + for (int32_t i = 0; i < actual_len; i++) { + float sum = 0; + for (int32_t n = 0; n < dur_shape_2; n++) { + sum += duration_processed[i * dur_shape_2 + n]; + } + sum /= speed; + pred_dur_original[i] = + static_cast(std::max(1.f, std::round(sum))); + } + + std::vector pred_dur_padding(max_seq_len_ - actual_len, 0); + pred_dur = pred_dur_original; + pred_dur.insert(pred_dur.end(), pred_dur_padding.begin(), + pred_dur_padding.end()); + + int32_t fixed_total_frames = max_seq_len_ * 2; + int32_t actual_frames = + std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); + int32_t diff = fixed_total_frames - actual_frames; + + if (diff < 0) { + auto indices = Argsort(pred_dur, actual_len, true); + int32_t decreased = 0; + for (auto idx : indices) { + if (pred_dur[idx] > 1 && decreased < std::abs(diff)) { + pred_dur[idx]--; + decreased++; + } + if (decreased >= std::abs(diff)) break; + } + } + + actual_frames = + std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); + int32_t remaining_frames = fixed_total_frames - actual_frames; + int32_t padding_len = max_seq_len_ - actual_len; + + if (remaining_frames > 0 && padding_len > 0) { + int32_t frames_per_padding = remaining_frames / padding_len; + int32_t remainder = remaining_frames % padding_len; + + for (int32_t i = actual_len; i < static_cast(pred_dur.size()); + i++) + pred_dur[i] = frames_per_padding; + + if (remainder > 0) { + for (int32_t i = actual_len; i < actual_len + remainder; i++) + pred_dur[i] += 1; + } + } + + total_frames = std::accumulate(pred_dur.begin(), pred_dur.end(), 0); + } + + std::vector CreateAlignmentMatrix(const std::vector &pred_dur, + int32_t total_frames) { + std::vector seq_range(max_seq_len_); + std::iota(seq_range.begin(), seq_range.end(), 0); + auto indices = NpRepeat(seq_range, pred_dur); + + std::vector pred_aln_trg(max_seq_len_ * total_frames, 0.0f); + if (!indices.empty()) { + int32_t col = 0; + for (auto i : indices) { + pred_aln_trg[i * total_frames + col] = 1.0f; + col++; + } + } + + return pred_aln_trg; + } + + void PostprocessXToAudio(const std::vector &x, int32_t num_frames, + std::vector &audio) { + int32_t n_fft = 20; + int32_t hop_length = 5; + int32_t half_n_fft = n_fft / 2 + 1; + + std::vector spec_part(half_n_fft * num_frames); + std::vector phase_part(half_n_fft * num_frames); + std::vector cos_part(half_n_fft * num_frames); + spec_part.assign(x.begin(), x.begin() + half_n_fft * num_frames); + phase_part.assign(x.begin() + half_n_fft * num_frames, x.end()); + + for (int32_t i = 0; i < half_n_fft * num_frames; i++) { + spec_part[i] = std::exp(spec_part[i]); + phase_part[i] = std::sin(phase_part[i]); + cos_part[i] = std::sqrt( + 1.f - std::max(0.f, std::min(phase_part[i] * phase_part[i], 1.0f))); + } + + knf::StftResult stft_result; + stft_result.num_frames = num_frames; + stft_result.real.resize(half_n_fft * num_frames); + stft_result.imag.resize(half_n_fft * num_frames); + + for (int32_t i = 0; i < half_n_fft; i++) { + for (int32_t n = 0; n < num_frames; n++) { + float spec = spec_part[i * num_frames + n]; + float real_part = spec * cos_part[i * num_frames + n]; + float imag_part = spec * phase_part[i * num_frames + n]; + + // StftResult layout is [frame_index * half_n_fft + bin] + stft_result.real[n * half_n_fft + i] = real_part; + stft_result.imag[n * half_n_fft + i] = imag_part; + } + } + + knf::StftConfig stft_config; + stft_config.n_fft = n_fft; + stft_config.hop_length = hop_length; + stft_config.win_length = n_fft; + stft_config.center = true; + stft_config.normalized = false; + stft_config.window_type = "hann"; + stft_config.pad_mode = "reflect"; + + knf::IStft istft(stft_config); + audio = istft.Compute(stft_result); + } + + void TrimAudioByContent(std::vector &audio, + int32_t actual_content_frames, int32_t total_frames, + int32_t actual_len) { + int32_t padding_len = max_seq_len_ - actual_len; + if (padding_len > 0) { + float content_ratio = actual_content_frames * 1.0f / total_frames; + int32_t audio_len_to_keep = + static_cast(audio.size() * content_ratio); + audio.resize(audio_len_to_keep); + } + } + + void PrepareInputIds(std::vector &input_ids, int32_t &actual_len, + bool &is_doubled) { + is_doubled = false; + int32_t original_actual_len = actual_len; + + if (actual_len <= 32) { // DOUBLE_INPUT_THRESHOLD + is_doubled = true; + std::vector valid_content(input_ids.begin(), + input_ids.begin() + actual_len); + std::vector input_ids_doubled; + input_ids_doubled.insert(input_ids_doubled.end(), valid_content.begin(), + valid_content.end()); + input_ids_doubled.insert(input_ids_doubled.end(), valid_content.begin(), + valid_content.end()); + + int32_t padding_len = max_seq_len_ - 2 * actual_len; + if (padding_len > 0) { + std::vector padding(padding_len, 0); + input_ids_doubled.insert(input_ids_doubled.end(), padding.begin(), + padding.end()); + } else { + input_ids_doubled.resize(max_seq_len_); + } + + input_ids = input_ids_doubled; + actual_len = std::min(original_actual_len * 2, max_seq_len_); + } + } + + void ComputeExternalPreprocessing(int32_t actual_len, + std::vector &text_mask) { + text_mask.resize(max_seq_len_); + for (int32_t i = 0; i < max_seq_len_; i++) { + text_mask[i] = (i >= actual_len) ? 1 : 0; + } + } + + private: + std::mutex mutex_; + AxEngineGuard ax_engine_guard_; + + OfflineTtsModelConfig config_; + OfflineTtsKokoroModelMetaData meta_data_; + + Ort::Env env_; + Ort::SessionOptions sess_opts_; + Ort::AllocatorWithDefaultOptions allocator_; + + std::unique_ptr model1_; + std::unique_ptr model2_; + std::unique_ptr model3_; + + std::unique_ptr model4_; + std::vector model4_input_names_; + std::vector model4_input_names_ptr_; + std::vector model4_output_names_; + std::vector model4_output_names_ptr_; +}; + +OfflineTtsKokoroModelAxera::~OfflineTtsKokoroModelAxera() = default; + +OfflineTtsKokoroModelAxera::OfflineTtsKokoroModelAxera( + const OfflineTtsModelConfig &config) + : impl_(std::make_unique(config)) {} + +template +OfflineTtsKokoroModelAxera::OfflineTtsKokoroModelAxera( + Manager *mgr, const OfflineTtsModelConfig &config) + : impl_(std::make_unique(mgr, config)) {} + +std::vector OfflineTtsKokoroModelAxera::Run( + const std::vector &x, int64_t sid, float speed) const { + return impl_->Run(x, sid, speed); +} + +const OfflineTtsKokoroModelMetaData &OfflineTtsKokoroModelAxera::GetMetaData() + const { + return impl_->GetMetaData(); +} + +#if __ANDROID_API__ >= 9 +template OfflineTtsKokoroModelAxera::OfflineTtsKokoroModelAxera( + AAssetManager *mgr, const OfflineTtsModelConfig &config); +#endif + +#if __OHOS__ +template OfflineTtsKokoroModelAxera::OfflineTtsKokoroModelAxera( + NativeResourceManager *mgr, const OfflineTtsModelConfig &config); +#endif + +} // namespace sherpa_onnx diff --git a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.h b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.h new file mode 100644 index 0000000000..97faaee02c --- /dev/null +++ b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.h @@ -0,0 +1,38 @@ +// sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.h +// +// Copyright (c) 2025 Xiaomi Corporation + +#ifndef SHERPA_ONNX_CSRC_AXERA_OFFLINE_TTS_KOKORO_MODEL_AXERA_H_ +#define SHERPA_ONNX_CSRC_AXERA_OFFLINE_TTS_KOKORO_MODEL_AXERA_H_ + +#include +#include +#include + +#include "sherpa-onnx/csrc/offline-tts-kokoro-model-meta-data.h" +#include "sherpa-onnx/csrc/offline-tts-model-config.h" + +namespace sherpa_onnx { + +class OfflineTtsKokoroModelAxera { + public: + ~OfflineTtsKokoroModelAxera(); + + explicit OfflineTtsKokoroModelAxera(const OfflineTtsModelConfig &config); + + template + OfflineTtsKokoroModelAxera(Manager *mgr, const OfflineTtsModelConfig &config); + + std::vector Run(const std::vector &x, int64_t sid = 0, + float speed = 1.0) const; + + const OfflineTtsKokoroModelMetaData &GetMetaData() const; + + private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace sherpa_onnx + +#endif // SHERPA_ONNX_CSRC_AXERA_OFFLINE_TTS_KOKORO_MODEL_AXERA_H_ diff --git a/sherpa-onnx/csrc/offline-tts-impl.cc b/sherpa-onnx/csrc/offline-tts-impl.cc index 84bffbd41e..1fcdd4b3c8 100644 --- a/sherpa-onnx/csrc/offline-tts-impl.cc +++ b/sherpa-onnx/csrc/offline-tts-impl.cc @@ -24,6 +24,10 @@ #include "sherpa-onnx/csrc/offline-tts-vits-impl.h" #include "sherpa-onnx/csrc/offline-tts-zipvoice-impl.h" +#ifdef SHERPA_ONNX_ENABLE_AXERA +#include "sherpa-onnx/csrc/axera/offline-tts-kokoro-axera-impl.h" +#endif + #ifdef SHERPA_ONNX_ENABLE_AXCL #include "sherpa-onnx/csrc/axcl/offline-tts-kokoro-axcl-impl.h" #endif @@ -44,11 +48,33 @@ std::vector OfflineTtsImpl::AddBlank(const std::vector &x, std::unique_ptr OfflineTtsImpl::Create( const OfflineTtsConfig &config) { + if (config.model.provider == "axera") { +#if SHERPA_ONNX_ENABLE_AXERA + if (!config.model.kokoro.model.empty()) { + return std::make_unique< + OfflineTtsKokoroAxeraImpl>(config); + } else { + SHERPA_ONNX_LOGE( + "Only Kokoro models are currently supported by axera for " + "non-streaming TTS."); + SHERPA_ONNX_EXIT(-1); + return nullptr; + } + +#else + SHERPA_ONNX_LOGE( + "Please rebuild sherpa-onnx with -DSHERPA_ONNX_ENABLE_AXERA=ON if you " + "want to use axera. "); + SHERPA_ONNX_EXIT(-1); + return nullptr; +#endif + } + if (config.model.provider == "axcl") { #if SHERPA_ONNX_ENABLE_AXCL if (!config.model.kokoro.model.empty()) { - return std::make_unique>( - config); + return std::make_unique< + OfflineTtsKokoroAxclImpl>(config); } else { SHERPA_ONNX_LOGE( "Only Kokoro models are currently supported by axcl for " @@ -92,6 +118,51 @@ std::unique_ptr OfflineTtsImpl::Create( template std::unique_ptr OfflineTtsImpl::Create( Manager *mgr, const OfflineTtsConfig &config) { + if (config.model.provider == "axera") { +#if SHERPA_ONNX_ENABLE_AXERA + if (!config.model.kokoro.model.empty()) { + return std::make_unique< + OfflineTtsKokoroAxeraImpl>(mgr, config); + } else { + SHERPA_ONNX_LOGE( + "Only Kokoro models are currently supported by axera for " + "non-streaming TTS."); + SHERPA_ONNX_EXIT(-1); + return nullptr; + } + +#else + SHERPA_ONNX_LOGE( + "Please rebuild sherpa-onnx with -DSHERPA_ONNX_ENABLE_AXERA=ON if you " + "want to use axera."); + SHERPA_ONNX_EXIT(-1); + return nullptr; +#endif + } + + if (config.model.provider == "axcl") { +#if SHERPA_ONNX_ENABLE_AXCL + if (!config.model.kokoro.model.empty()) { + return std::make_unique< + OfflineTtsKokoroAxclImpl>(mgr, config); + } else { + SHERPA_ONNX_LOGE( + "Only Kokoro models are currently supported by axcl for " + "non-streaming TTS."); + SHERPA_ONNX_EXIT(-1); + return nullptr; + } + +#else + SHERPA_ONNX_LOGE( + "Please rebuild sherpa-onnx with -DSHERPA_ONNX_ENABLE_AXCL=ON if you " + "want to use axcl. See also " + "https://k2-fsa.github.io/sherpa/onnx/axcl/install.html"); + SHERPA_ONNX_EXIT(-1); + return nullptr; +#endif + } + if (!config.model.vits.model.empty()) { return std::make_unique(mgr, config); } else if (!config.model.matcha.acoustic_model.empty()) { From dd293066a2a20d0fd74cd83877bd9ce37282b42c Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Wed, 11 Mar 2026 15:42:15 +0800 Subject: [PATCH 04/10] Read metadata from model4 for AXCL and AXERA --- .../axcl/offline-tts-kokoro-model-axcl.cc | 137 ++++++++++++------ .../axera/offline-tts-kokoro-model-axera.cc | 57 ++++++-- 2 files changed, 132 insertions(+), 62 deletions(-) diff --git a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc index fd5f65deed..ff431cca82 100644 --- a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc +++ b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc @@ -30,6 +30,7 @@ #include "sherpa-onnx/csrc/macros.h" #include "sherpa-onnx/csrc/onnx-utils.h" #include "sherpa-onnx/csrc/session.h" +#include "sherpa-onnx/csrc/text-utils.h" namespace sherpa_onnx { @@ -111,7 +112,8 @@ class OfflineTtsKokoroModelAxcl::Impl { return meta_data_; } - std::vector Run(std::vector input_ids, int64_t sid, float speed) { + std::vector Run(std::vector input_ids, int64_t sid, + float speed) { if (input_ids.size() < 2) { return {}; } @@ -138,15 +140,15 @@ class OfflineTtsKokoroModelAxcl::Impl { private: void Init(const OfflineTtsModelConfig &config) { std::string model_path = config.kokoro.model; - - // kokoro axcl mode we assume config.kokoro.model is a directory containing models - // or we can use the pattern like kokoro.model + "_part1.axmodel" - // Let's assume the user passes a directory. Or the user passes the prefix path. - // For simplicity, we assume `config.kokoro.model` is a directory. + + // kokoro axcl mode we assume config.kokoro.model is a directory containing + // models or we can use the pattern like kokoro.model + "_part1.axmodel" + // Let's assume the user passes a directory. Or the user passes the prefix + // path. For simplicity, we assume `config.kokoro.model` is a directory. std::string model1_path = model_path + "/kokoro_part1.axmodel"; std::string model2_path = model_path + "/kokoro_part2.axmodel"; std::string model3_path = model_path + "/kokoro_part3.axmodel"; - std::string model4_path = model_path + "/kokoro_part4.onnx"; + std::string model4_path = model_path + "/kokoro_part4.onnx"; model1_ = std::make_unique(model1_path); model2_ = std::make_unique(model2_path); @@ -155,7 +157,8 @@ class OfflineTtsKokoroModelAxcl::Impl { auto buf4 = ReadFile(model4_path); model4_ = std::make_unique(env_, buf4.data(), buf4.size(), sess_opts_); - GetInputNames(model4_.get(), &model4_input_names_, &model4_input_names_ptr_); + GetInputNames(model4_.get(), &model4_input_names_, + &model4_input_names_ptr_); GetOutputNames(model4_.get(), &model4_output_names_, &model4_output_names_ptr_); @@ -169,7 +172,7 @@ class OfflineTtsKokoroModelAxcl::Impl { std::string model1_path = model_path + "/kokoro_part1.axmodel"; std::string model2_path = model_path + "/kokoro_part2.axmodel"; std::string model3_path = model_path + "/kokoro_part3.axmodel"; - std::string model4_path = model_path + "/kokoro_part4.onnx"; + std::string model4_path = model_path + "/kokoro_part4.onnx"; auto buf1 = ReadFile(mgr, model1_path); model1_ = std::make_unique(buf1.data(), buf1.size()); @@ -183,7 +186,8 @@ class OfflineTtsKokoroModelAxcl::Impl { auto buf4 = ReadFile(mgr, model4_path); model4_ = std::make_unique(env_, buf4.data(), buf4.size(), sess_opts_); - GetInputNames(model4_.get(), &model4_input_names_, &model4_input_names_ptr_); + GetInputNames(model4_.get(), &model4_input_names_, + &model4_input_names_ptr_); GetOutputNames(model4_.get(), &model4_output_names_, &model4_output_names_ptr_); @@ -192,26 +196,53 @@ class OfflineTtsKokoroModelAxcl::Impl { } void LoadVoices(const char *voices_data, size_t voices_data_length) { - // get metadata - meta_data_.sample_rate = 24000; - meta_data_.num_speakers = voices_data_length / (sizeof(float) * 256 * 510); - meta_data_.version = 1; - meta_data_.has_espeak = 1; - meta_data_.voice = "en-us"; - meta_data_.max_token_len = 510; + // Get metadata from model4 (ONNX model) + Ort::ModelMetadata meta_data = model4_->GetModelMetadata(); + Ort::AllocatorWithDefaultOptions allocator; + + SHERPA_ONNX_READ_META_DATA(meta_data_.sample_rate, "sample_rate"); + SHERPA_ONNX_READ_META_DATA_WITH_DEFAULT(meta_data_.version, "version", 1); + SHERPA_ONNX_READ_META_DATA(meta_data_.num_speakers, "n_speakers"); + SHERPA_ONNX_READ_META_DATA(meta_data_.has_espeak, "has_espeak"); + SHERPA_ONNX_READ_META_DATA_STR_WITH_DEFAULT(meta_data_.voice, "voice", + "en-us"); + SHERPA_ONNX_READ_META_DATA_VEC(style_dim_, "style_dim"); + + if (style_dim_.size() != 3) { + SHERPA_ONNX_LOGE("style_dim should be 3-d, given: %d", + static_cast(style_dim_.size())); + SHERPA_ONNX_EXIT(-1); + } + + if (style_dim_[1] != 1) { + SHERPA_ONNX_LOGE("style_dim[1] should be 1, given: %d", style_dim_[1]); + SHERPA_ONNX_EXIT(-1); + } + + meta_data_.max_token_len = style_dim_[0]; int32_t actual_num_floats = voices_data_length / sizeof(float); + int32_t expected_num_floats = + style_dim_[0] * style_dim_[2] * meta_data_.num_speakers; + + if (actual_num_floats != expected_num_floats) { + SHERPA_ONNX_LOGE( + "Corrupted voices file. Expected #floats: %d, actual: %d", + expected_num_floats, actual_num_floats); + SHERPA_ONNX_EXIT(-1); + } + styles_ = std::vector( reinterpret_cast(voices_data), reinterpret_cast(voices_data) + actual_num_floats); - + auto shape1 = model1_->TensorShape("input_ids"); max_seq_len_ = shape1[1]; } std::vector LoadVoiceEmbedding(int32_t sid, int32_t phoneme_len) { - int32_t style_dim0 = 510; // max_token_len - int32_t style_dim1 = 256; + int32_t style_dim0 = style_dim_[0]; + int32_t style_dim1 = style_dim_[2]; phoneme_len = std::max(phoneme_len, 0); std::vector ref_s(style_dim1); @@ -233,8 +264,9 @@ class OfflineTtsKokoroModelAxcl::Impl { std::vector &audio) { int32_t actual_len = input_ids.size(); if (actual_len > max_seq_len_) { - SHERPA_ONNX_LOGE("Input ids length %d exceeds max_seq_len %d, truncating.", - actual_len, max_seq_len_); + SHERPA_ONNX_LOGE( + "Input ids length %d exceeds max_seq_len %d, truncating.", actual_len, + max_seq_len_); input_ids.resize(max_seq_len_); actual_len = max_seq_len_; } @@ -279,7 +311,8 @@ class OfflineTtsKokoroModelAxcl::Impl { model1_->SetInputTensorData("input_ids", input_ids_i32.data(), input_ids_i32.size()); model1_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); - model1_->SetInputTensorData("text_mask", text_mask_u8.data(), text_mask_u8.size()); + model1_->SetInputTensorData("text_mask", text_mask_u8.data(), + text_mask_u8.size()); if (!model1_->Run()) { return false; @@ -291,10 +324,12 @@ class OfflineTtsKokoroModelAxcl::Impl { std::vector pred_dur; ProcessDuration(duration, actual_len, speed, pred_dur, total_frames); - std::vector pred_aln_trg = CreateAlignmentMatrix(pred_dur, total_frames); + std::vector pred_aln_trg = + CreateAlignmentMatrix(pred_dur, total_frames); - int32_t d_shape_1 = max_seq_len_; - int32_t d_shape_2 = 640; // style_dim_[2] essentially, or based on d_shape_[2] + int32_t d_shape_1 = max_seq_len_; + int32_t d_shape_2 = + 640; // style_dim_[2] essentially, or based on d_shape_[2] std::vector en(total_frames * d_shape_2, 0.0f); for (int32_t j = 0; j < d_shape_2; ++j) { @@ -311,9 +346,12 @@ class OfflineTtsKokoroModelAxcl::Impl { model2_->SetInputTensorData("en", en.data(), en.size()); model2_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); - model2_->SetInputTensorData("input_ids", input_ids_i32.data(), input_ids_i32.size()); - model2_->SetInputTensorData("text_mask", text_mask_float.data(), text_mask_float.size()); - model2_->SetInputTensorData("pred_aln_trg", pred_aln_trg.data(), pred_aln_trg.size()); + model2_->SetInputTensorData("input_ids", input_ids_i32.data(), + input_ids_i32.size()); + model2_->SetInputTensorData("text_mask", text_mask_float.data(), + text_mask_float.size()); + model2_->SetInputTensorData("pred_aln_trg", pred_aln_trg.data(), + pred_aln_trg.size()); if (!model2_->Run()) { return false; @@ -326,9 +364,9 @@ class OfflineTtsKokoroModelAxcl::Impl { auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); std::array f0_shape = {1, static_cast(f0_pred.size())}; - Ort::Value f0_tensor = Ort::Value::CreateTensor( - memory_info, f0_pred.data(), f0_pred.size(), f0_shape.data(), - f0_shape.size()); + Ort::Value f0_tensor = + Ort::Value::CreateTensor(memory_info, f0_pred.data(), f0_pred.size(), + f0_shape.data(), f0_shape.size()); auto out4 = model4_->Run({}, model4_input_names_ptr_.data(), &f0_tensor, 1, model4_output_names_ptr_.data(), @@ -355,8 +393,8 @@ class OfflineTtsKokoroModelAxcl::Impl { PostprocessXToAudio(x_out, 23041, audio); if (is_doubled) { - actual_content_frames = std::accumulate(pred_dur.begin(), - pred_dur.begin() + original_actual_len, 0); + actual_content_frames = std::accumulate( + pred_dur.begin(), pred_dur.begin() + original_actual_len, 0); audio.erase(audio.begin() + audio.size() / 2, audio.end()); total_frames = total_frames / 2; } else { @@ -371,15 +409,16 @@ class OfflineTtsKokoroModelAxcl::Impl { int32_t &total_frames) { std::vector pred_dur_original(actual_len, 0); std::vector duration_processed = Sigmoid(duration); - - int32_t dur_shape_2 = 50; + + int32_t dur_shape_2 = 50; for (int32_t i = 0; i < actual_len; i++) { float sum = 0; for (int32_t n = 0; n < dur_shape_2; n++) { sum += duration_processed[i * dur_shape_2 + n]; } sum /= speed; - pred_dur_original[i] = static_cast(std::max(1.f, std::round(sum))); + pred_dur_original[i] = + static_cast(std::max(1.f, std::round(sum))); } std::vector pred_dur_padding(max_seq_len_ - actual_len, 0); @@ -413,7 +452,8 @@ class OfflineTtsKokoroModelAxcl::Impl { int32_t frames_per_padding = remaining_frames / padding_len; int32_t remainder = remaining_frames % padding_len; - for (int32_t i = actual_len; i < static_cast(pred_dur.size()); i++) + for (int32_t i = actual_len; i < static_cast(pred_dur.size()); + i++) pred_dur[i] = frames_per_padding; if (remainder > 0) { @@ -458,8 +498,8 @@ class OfflineTtsKokoroModelAxcl::Impl { for (int32_t i = 0; i < half_n_fft * num_frames; i++) { spec_part[i] = std::exp(spec_part[i]); phase_part[i] = std::sin(phase_part[i]); - cos_part[i] = - std::sqrt(1.f - std::max(0.f, std::min(phase_part[i] * phase_part[i], 1.0f))); + cos_part[i] = std::sqrt( + 1.f - std::max(0.f, std::min(phase_part[i] * phase_part[i], 1.0f))); } knf::StftResult stft_result; @@ -492,12 +532,14 @@ class OfflineTtsKokoroModelAxcl::Impl { audio = istft.Compute(stft_result); } - void TrimAudioByContent(std::vector &audio, int32_t actual_content_frames, - int32_t total_frames, int32_t actual_len) { + void TrimAudioByContent(std::vector &audio, + int32_t actual_content_frames, int32_t total_frames, + int32_t actual_len) { int32_t padding_len = max_seq_len_ - actual_len; if (padding_len > 0) { float content_ratio = actual_content_frames * 1.0f / total_frames; - int32_t audio_len_to_keep = static_cast(audio.size() * content_ratio); + int32_t audio_len_to_keep = + static_cast(audio.size() * content_ratio); audio.resize(audio_len_to_keep); } } @@ -507,7 +549,7 @@ class OfflineTtsKokoroModelAxcl::Impl { is_doubled = false; int32_t original_actual_len = actual_len; - if (actual_len <= 32) { // DOUBLE_INPUT_THRESHOLD + if (actual_len <= 32) { // DOUBLE_INPUT_THRESHOLD is_doubled = true; std::vector valid_content(input_ids.begin(), input_ids.begin() + actual_len); @@ -545,6 +587,7 @@ class OfflineTtsKokoroModelAxcl::Impl { private: OfflineTtsModelConfig config_; OfflineTtsKokoroModelMetaData meta_data_; + std::vector style_dim_; Ort::Env env_; Ort::SessionOptions sess_opts_; @@ -572,8 +615,9 @@ OfflineTtsKokoroModelAxcl::OfflineTtsKokoroModelAxcl( Manager *mgr, const OfflineTtsModelConfig &config) : impl_(std::make_unique(mgr, config)) {} -std::vector OfflineTtsKokoroModelAxcl::Run( - const std::vector& x, int64_t sid, float speed) const { +std::vector OfflineTtsKokoroModelAxcl::Run(const std::vector &x, + int64_t sid, + float speed) const { return impl_->Run(x, sid, speed); } @@ -593,4 +637,3 @@ template OfflineTtsKokoroModelAxcl::OfflineTtsKokoroModelAxcl( #endif } // namespace sherpa_onnx - diff --git a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc index 7699f1aedd..be23d19a4d 100644 --- a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc +++ b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc @@ -35,6 +35,7 @@ #include "sherpa-onnx/csrc/macros.h" #include "sherpa-onnx/csrc/onnx-utils.h" #include "sherpa-onnx/csrc/session.h" +#include "sherpa-onnx/csrc/text-utils.h" namespace sherpa_onnx { @@ -213,10 +214,6 @@ static std::string GetModelPath(const std::string &model, const std::string &p, } class OfflineTtsKokoroModelAxera::Impl { - private: - int32_t max_seq_len_ = 0; - std::vector styles_; - public: explicit Impl(const OfflineTtsModelConfig &config) : config_(config), @@ -319,15 +316,41 @@ class OfflineTtsKokoroModelAxera::Impl { } void LoadVoices(const char *voices_data, size_t voices_data_length) { - // get metadata - meta_data_.sample_rate = 24000; - meta_data_.num_speakers = voices_data_length / (sizeof(float) * 256 * 510); - meta_data_.version = 1; - meta_data_.has_espeak = 1; - meta_data_.voice = "en-us"; - meta_data_.max_token_len = 510; + // Get metadata from model4 (ONNX model) + Ort::ModelMetadata meta_data = model4_->GetModelMetadata(); + Ort::AllocatorWithDefaultOptions allocator; + + SHERPA_ONNX_READ_META_DATA(meta_data_.sample_rate, "sample_rate"); + SHERPA_ONNX_READ_META_DATA_WITH_DEFAULT(meta_data_.version, "version", 1); + SHERPA_ONNX_READ_META_DATA(meta_data_.num_speakers, "n_speakers"); + SHERPA_ONNX_READ_META_DATA(meta_data_.has_espeak, "has_espeak"); + SHERPA_ONNX_READ_META_DATA_STR_WITH_DEFAULT(meta_data_.voice, "voice", + "en-us"); + SHERPA_ONNX_READ_META_DATA_VEC(style_dim_, "style_dim"); + + if (style_dim_.size() != 3) { + SHERPA_ONNX_LOGE("style_dim should be 3-d, given: %d", + static_cast(style_dim_.size())); + SHERPA_ONNX_EXIT(-1); + } + + if (style_dim_[1] != 1) { + SHERPA_ONNX_LOGE("style_dim[1] should be 1, given: %d", style_dim_[1]); + SHERPA_ONNX_EXIT(-1); + } + + meta_data_.max_token_len = style_dim_[0]; int32_t actual_num_floats = voices_data_length / sizeof(float); + int32_t expected_num_floats = + style_dim_[0] * style_dim_[2] * meta_data_.num_speakers; + + if (actual_num_floats != expected_num_floats) { + SHERPA_ONNX_LOGE( + "Corrupted voices file. Expected #floats: %d, actual: %d", + expected_num_floats, actual_num_floats); + SHERPA_ONNX_EXIT(-1); + } styles_ = std::vector( reinterpret_cast(voices_data), reinterpret_cast(voices_data) + actual_num_floats); @@ -343,8 +366,8 @@ class OfflineTtsKokoroModelAxera::Impl { } std::vector LoadVoiceEmbedding(int32_t sid, int32_t phoneme_len) { - int32_t style_dim0 = 510; - int32_t style_dim1 = 256; + int32_t style_dim0 = style_dim_[0]; + int32_t style_dim1 = style_dim_[2]; phoneme_len = std::max(phoneme_len, 0); sid = std::max(sid, 0); @@ -696,8 +719,6 @@ class OfflineTtsKokoroModelAxera::Impl { AxEngineGuard ax_engine_guard_; OfflineTtsModelConfig config_; - OfflineTtsKokoroModelMetaData meta_data_; - Ort::Env env_; Ort::SessionOptions sess_opts_; Ort::AllocatorWithDefaultOptions allocator_; @@ -711,6 +732,12 @@ class OfflineTtsKokoroModelAxera::Impl { std::vector model4_input_names_ptr_; std::vector model4_output_names_; std::vector model4_output_names_ptr_; + + OfflineTtsKokoroModelMetaData meta_data_; + std::vector style_dim_; + std::vector styles_; + + int32_t max_seq_len_; }; OfflineTtsKokoroModelAxera::~OfflineTtsKokoroModelAxera() = default; From 4cdc81fef44f203562f244aec22b68ec0dc0960e Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Wed, 11 Mar 2026 17:54:35 +0800 Subject: [PATCH 05/10] use comma-separated model paths instead of directory --- .../axcl/offline-tts-kokoro-model-axcl.cc | 158 +++++++++++------- .../axera/offline-tts-kokoro-model-axera.cc | 152 +++++++++++------ .../csrc/offline-tts-kokoro-model-config.cc | 12 +- 3 files changed, 209 insertions(+), 113 deletions(-) diff --git a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc index ff431cca82..f415e7a679 100644 --- a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc +++ b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc @@ -139,22 +139,45 @@ class OfflineTtsKokoroModelAxcl::Impl { private: void Init(const OfflineTtsModelConfig &config) { - std::string model_path = config.kokoro.model; - - // kokoro axcl mode we assume config.kokoro.model is a directory containing - // models or we can use the pattern like kokoro.model + "_part1.axmodel" - // Let's assume the user passes a directory. Or the user passes the prefix - // path. For simplicity, we assume `config.kokoro.model` is a directory. - std::string model1_path = model_path + "/kokoro_part1.axmodel"; - std::string model2_path = model_path + "/kokoro_part2.axmodel"; - std::string model3_path = model_path + "/kokoro_part3.axmodel"; - std::string model4_path = model_path + "/kokoro_part4.onnx"; - - model1_ = std::make_unique(model1_path); - model2_ = std::make_unique(model2_path); - model3_ = std::make_unique(model3_path); - - auto buf4 = ReadFile(model4_path); + std::string model_str = config.kokoro.model; + + std::vector model_files; + SplitStringToVector(model_str, ",", false, &model_files); + + std::vector model_paths; + + if (config_.provider == "axcl") { + // axcl provider requires exactly 4 model files + if (model_files.size() != 4) { + SHERPA_ONNX_LOGE( + "For axcl provider, model should contain exactly 4 files " + "separated by comma. Given %d files: %s", + static_cast(model_files.size()), model_str.c_str()); + SHERPA_ONNX_EXIT(-1); + } + + // Validate all 4 files exist + for (int32_t i = 0; i < 4; ++i) { + if (!FileExists(model_files[i])) { + SHERPA_ONNX_LOGE("Model file '%s' does not exist", + model_files[i].c_str()); + SHERPA_ONNX_EXIT(-1); + } + } + model_paths = model_files; + } else { + SHERPA_ONNX_LOGE( + "This model only supports axcl provider. Please use provider=axcl"); + SHERPA_ONNX_EXIT(-1); + } + + // Load ax models (indices 0, 1, 2) + for (int32_t i = 0; i < 3; ++i) { + axcl_models_.push_back(std::make_unique(model_paths[i])); + } + + // Load onnx model (index 3) + auto buf4 = ReadFile(model_paths[3]); model4_ = std::make_unique(env_, buf4.data(), buf4.size(), sess_opts_); GetInputNames(model4_.get(), &model4_input_names_, @@ -168,22 +191,38 @@ class OfflineTtsKokoroModelAxcl::Impl { template void Init(Manager *mgr, const OfflineTtsModelConfig &config) { - std::string model_path = config.kokoro.model; - std::string model1_path = model_path + "/kokoro_part1.axmodel"; - std::string model2_path = model_path + "/kokoro_part2.axmodel"; - std::string model3_path = model_path + "/kokoro_part3.axmodel"; - std::string model4_path = model_path + "/kokoro_part4.onnx"; + std::string model_str = config.kokoro.model; - auto buf1 = ReadFile(mgr, model1_path); - model1_ = std::make_unique(buf1.data(), buf1.size()); + std::vector model_files; + SplitStringToVector(model_str, ",", false, &model_files); - auto buf2 = ReadFile(mgr, model2_path); - model2_ = std::make_unique(buf2.data(), buf2.size()); + std::vector model_paths; + + if (config_.provider == "axcl") { + // axcl provider requires exactly 4 model files + if (model_files.size() != 4) { + SHERPA_ONNX_LOGE( + "For axcl provider, model should contain exactly 4 files " + "separated by comma. Given %d files: %s", + static_cast(model_files.size()), model_str.c_str()); + SHERPA_ONNX_EXIT(-1); + } + model_paths = model_files; + } else { + SHERPA_ONNX_LOGE( + "This model only supports axcl provider. Please use provider=axcl"); + SHERPA_ONNX_EXIT(-1); + } - auto buf3 = ReadFile(mgr, model3_path); - model3_ = std::make_unique(buf3.data(), buf3.size()); + // Load ax models (indices 0, 1, 2) + for (int32_t i = 0; i < 3; ++i) { + auto buf = ReadFile(mgr, model_paths[i]); + axcl_models_.push_back( + std::make_unique(buf.data(), buf.size())); + } - auto buf4 = ReadFile(mgr, model4_path); + // Load onnx model (index 3) + auto buf4 = ReadFile(mgr, model_paths[3]); model4_ = std::make_unique(env_, buf4.data(), buf4.size(), sess_opts_); GetInputNames(model4_.get(), &model4_input_names_, @@ -236,7 +275,7 @@ class OfflineTtsKokoroModelAxcl::Impl { reinterpret_cast(voices_data), reinterpret_cast(voices_data) + actual_num_floats); - auto shape1 = model1_->TensorShape("input_ids"); + auto shape1 = axcl_models_[0]->TensorShape("input_ids"); max_seq_len_ = shape1[1]; } @@ -308,18 +347,19 @@ class OfflineTtsKokoroModelAxcl::Impl { std::vector text_mask_u8(text_mask.begin(), text_mask.end()); - model1_->SetInputTensorData("input_ids", input_ids_i32.data(), - input_ids_i32.size()); - model1_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); - model1_->SetInputTensorData("text_mask", text_mask_u8.data(), - text_mask_u8.size()); + axcl_models_[0]->SetInputTensorData("input_ids", input_ids_i32.data(), + input_ids_i32.size()); + axcl_models_[0]->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); + axcl_models_[0]->SetInputTensorData("text_mask", text_mask_u8.data(), + text_mask_u8.size()); - if (!model1_->Run()) { + if (!axcl_models_[0]->Run()) { return false; } - std::vector duration = model1_->GetOutputTensorData("duration"); - std::vector d = model1_->GetOutputTensorData("d"); + std::vector duration = + axcl_models_[0]->GetOutputTensorData("duration"); + std::vector d = axcl_models_[0]->GetOutputTensorData("d"); std::vector pred_dur; ProcessDuration(duration, actual_len, speed, pred_dur, total_frames); @@ -344,22 +384,23 @@ class OfflineTtsKokoroModelAxcl::Impl { std::vector text_mask_float(text_mask.begin(), text_mask.end()); - model2_->SetInputTensorData("en", en.data(), en.size()); - model2_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); - model2_->SetInputTensorData("input_ids", input_ids_i32.data(), - input_ids_i32.size()); - model2_->SetInputTensorData("text_mask", text_mask_float.data(), - text_mask_float.size()); - model2_->SetInputTensorData("pred_aln_trg", pred_aln_trg.data(), - pred_aln_trg.size()); + axcl_models_[1]->SetInputTensorData("en", en.data(), en.size()); + axcl_models_[1]->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); + axcl_models_[1]->SetInputTensorData("input_ids", input_ids_i32.data(), + input_ids_i32.size()); + axcl_models_[1]->SetInputTensorData("text_mask", text_mask_float.data(), + text_mask_float.size()); + axcl_models_[1]->SetInputTensorData("pred_aln_trg", pred_aln_trg.data(), + pred_aln_trg.size()); - if (!model2_->Run()) { + if (!axcl_models_[1]->Run()) { return false; } - std::vector f0_pred = model2_->GetOutputTensorData("F0_pred"); - std::vector n_pred = model2_->GetOutputTensorData("N_pred"); - std::vector asr = model2_->GetOutputTensorData("asr"); + std::vector f0_pred = + axcl_models_[1]->GetOutputTensorData("F0_pred"); + std::vector n_pred = axcl_models_[1]->GetOutputTensorData("N_pred"); + std::vector asr = axcl_models_[1]->GetOutputTensorData("asr"); auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); @@ -379,17 +420,18 @@ class OfflineTtsKokoroModelAxcl::Impl { } std::vector har(p_har, p_har + har_num_elements); - model3_->SetInputTensorData("asr", asr.data(), asr.size()); - model3_->SetInputTensorData("F0_pred", f0_pred.data(), f0_pred.size()); - model3_->SetInputTensorData("N_pred", n_pred.data(), n_pred.size()); - model3_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); - model3_->SetInputTensorData("har", har.data(), har.size()); + axcl_models_[2]->SetInputTensorData("asr", asr.data(), asr.size()); + axcl_models_[2]->SetInputTensorData("F0_pred", f0_pred.data(), + f0_pred.size()); + axcl_models_[2]->SetInputTensorData("N_pred", n_pred.data(), n_pred.size()); + axcl_models_[2]->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); + axcl_models_[2]->SetInputTensorData("har", har.data(), har.size()); - if (!model3_->Run()) { + if (!axcl_models_[2]->Run()) { return false; } - std::vector x_out = model3_->GetOutputTensorData("x"); + std::vector x_out = axcl_models_[2]->GetOutputTensorData("x"); PostprocessXToAudio(x_out, 23041, audio); if (is_doubled) { @@ -593,9 +635,7 @@ class OfflineTtsKokoroModelAxcl::Impl { Ort::SessionOptions sess_opts_; Ort::AllocatorWithDefaultOptions allocator_; - std::unique_ptr model1_; - std::unique_ptr model2_; - std::unique_ptr model3_; + std::vector> axcl_models_; std::unique_ptr model4_; std::vector model4_input_names_; diff --git a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc index be23d19a4d..a0fe960d09 100644 --- a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc +++ b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc @@ -266,16 +266,45 @@ class OfflineTtsKokoroModelAxera::Impl { private: void Init(const OfflineTtsModelConfig &config) { - std::string model1_path = GetModelPath(config.kokoro.model, "part1"); - std::string model2_path = GetModelPath(config.kokoro.model, "part2"); - std::string model3_path = GetModelPath(config.kokoro.model, "part3"); - std::string model4_path = GetModelPath(config.kokoro.model, "part4", true); + std::string model_str = config.kokoro.model; - model1_ = std::make_unique(model1_path); - model2_ = std::make_unique(model2_path); - model3_ = std::make_unique(model3_path); + std::vector model_files; + SplitStringToVector(model_str, ",", false, &model_files); - auto buf4 = ReadFile(model4_path); + std::vector model_paths; + + if (config_.provider == "axera") { + // axera provider requires exactly 4 model files + if (model_files.size() != 4) { + SHERPA_ONNX_LOGE( + "For axera provider, model should contain exactly 4 files " + "separated by comma. Given %d files: %s", + static_cast(model_files.size()), model_str.c_str()); + SHERPA_ONNX_EXIT(-1); + } + + // Validate all 4 files exist + for (int32_t i = 0; i < 4; ++i) { + if (!FileExists(model_files[i])) { + SHERPA_ONNX_LOGE("Model file '%s' does not exist", + model_files[i].c_str()); + SHERPA_ONNX_EXIT(-1); + } + } + model_paths = model_files; + } else { + SHERPA_ONNX_LOGE( + "This model only supports axera provider. Please use provider=axera"); + SHERPA_ONNX_EXIT(-1); + } + + // Load axera models (indices 0, 1, 2) + for (int32_t i = 0; i < 3; ++i) { + axera_models_.push_back(std::make_unique(model_paths[i])); + } + + // Load onnx model (index 3) + auto buf4 = ReadFile(model_paths[3]); model4_ = std::make_unique(env_, buf4.data(), buf4.size(), sess_opts_); GetInputNames(model4_.get(), &model4_input_names_, @@ -289,21 +318,38 @@ class OfflineTtsKokoroModelAxera::Impl { template void Init(Manager *mgr, const OfflineTtsModelConfig &config) { - std::string model1_path = GetModelPath(config.kokoro.model, "part1"); - std::string model2_path = GetModelPath(config.kokoro.model, "part2"); - std::string model3_path = GetModelPath(config.kokoro.model, "part3"); - std::string model4_path = GetModelPath(config.kokoro.model, "part4", true); + std::string model_str = config.kokoro.model; - auto buf1 = ReadFile(mgr, model1_path); - model1_ = std::make_unique(buf1.data(), buf1.size()); + std::vector model_files; + SplitStringToVector(model_str, ",", false, &model_files); - auto buf2 = ReadFile(mgr, model2_path); - model2_ = std::make_unique(buf2.data(), buf2.size()); + std::vector model_paths; - auto buf3 = ReadFile(mgr, model3_path); - model3_ = std::make_unique(buf3.data(), buf3.size()); + if (config_.provider == "axera") { + // axera provider requires exactly 4 model files + if (model_files.size() != 4) { + SHERPA_ONNX_LOGE( + "For axera provider, model should contain exactly 4 files " + "separated by comma. Given %d files: %s", + static_cast(model_files.size()), model_str.c_str()); + SHERPA_ONNX_EXIT(-1); + } + model_paths = model_files; + } else { + SHERPA_ONNX_LOGE( + "This model only supports axera provider. Please use provider=axera"); + SHERPA_ONNX_EXIT(-1); + } + + // Load axera models (indices 0, 1, 2) + for (int32_t i = 0; i < 3; ++i) { + auto buf = ReadFile(mgr, model_paths[i]); + axera_models_.push_back( + std::make_unique(buf.data(), buf.size())); + } - auto buf4 = ReadFile(mgr, model4_path); + // Load onnx model (index 3) + auto buf4 = ReadFile(mgr, model_paths[3]); model4_ = std::make_unique(env_, buf4.data(), buf4.size(), sess_opts_); GetInputNames(model4_.get(), &model4_input_names_, @@ -355,7 +401,7 @@ class OfflineTtsKokoroModelAxera::Impl { reinterpret_cast(voices_data), reinterpret_cast(voices_data) + actual_num_floats); - auto shape1 = model1_->TensorShape("input_ids"); + auto shape1 = axera_models_[0]->TensorShape("input_ids"); if (shape1.size() < 2) { SHERPA_ONNX_LOGE("Unexpected shape rank for input_ids: %d", static_cast(shape1.size())); @@ -436,20 +482,22 @@ class OfflineTtsKokoroModelAxera::Impl { std::vector text_mask_u8(text_mask.begin(), text_mask.end()); - if (!model1_->SetInputTensorData("input_ids", input_ids_i32.data(), - input_ids_i32.size()) || - !model1_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()) || - !model1_->SetInputTensorData("text_mask", text_mask_u8.data(), - text_mask_u8.size())) { + if (!axera_models_[0]->SetInputTensorData("input_ids", input_ids_i32.data(), + input_ids_i32.size()) || + !axera_models_[0]->SetInputTensorData("ref_s", ref_s.data(), + ref_s.size()) || + !axera_models_[0]->SetInputTensorData("text_mask", text_mask_u8.data(), + text_mask_u8.size())) { return false; } - if (!model1_->Run()) { + if (!axera_models_[0]->Run()) { return false; } - std::vector duration = model1_->GetOutputTensorData("duration"); - std::vector d = model1_->GetOutputTensorData("d"); + std::vector duration = + axera_models_[0]->GetOutputTensorData("duration"); + std::vector d = axera_models_[0]->GetOutputTensorData("d"); std::vector pred_dur; ProcessDuration(duration, actual_len, speed, pred_dur, total_frames); @@ -472,24 +520,26 @@ class OfflineTtsKokoroModelAxera::Impl { std::vector text_mask_float(text_mask.begin(), text_mask.end()); - if (!model2_->SetInputTensorData("en", en.data(), en.size()) || - !model2_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()) || - !model2_->SetInputTensorData("input_ids", input_ids_i32.data(), - input_ids_i32.size()) || - !model2_->SetInputTensorData("text_mask", text_mask_float.data(), - text_mask_float.size()) || - !model2_->SetInputTensorData("pred_aln_trg", pred_aln_trg.data(), - pred_aln_trg.size())) { + if (!axera_models_[1]->SetInputTensorData("en", en.data(), en.size()) || + !axera_models_[1]->SetInputTensorData("ref_s", ref_s.data(), + ref_s.size()) || + !axera_models_[1]->SetInputTensorData("input_ids", input_ids_i32.data(), + input_ids_i32.size()) || + !axera_models_[1]->SetInputTensorData( + "text_mask", text_mask_float.data(), text_mask_float.size()) || + !axera_models_[1]->SetInputTensorData( + "pred_aln_trg", pred_aln_trg.data(), pred_aln_trg.size())) { return false; } - if (!model2_->Run()) { + if (!axera_models_[1]->Run()) { return false; } - std::vector f0_pred = model2_->GetOutputTensorData("F0_pred"); - std::vector n_pred = model2_->GetOutputTensorData("N_pred"); - std::vector asr = model2_->GetOutputTensorData("asr"); + std::vector f0_pred = + axera_models_[1]->GetOutputTensorData("F0_pred"); + std::vector n_pred = axera_models_[1]->GetOutputTensorData("N_pred"); + std::vector asr = axera_models_[1]->GetOutputTensorData("asr"); auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); @@ -509,20 +559,22 @@ class OfflineTtsKokoroModelAxera::Impl { } std::vector har(p_har, p_har + har_num_elements); - if (!model3_->SetInputTensorData("asr", asr.data(), asr.size()) || - !model3_->SetInputTensorData("F0_pred", f0_pred.data(), - f0_pred.size()) || - !model3_->SetInputTensorData("N_pred", n_pred.data(), n_pred.size()) || - !model3_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()) || - !model3_->SetInputTensorData("har", har.data(), har.size())) { + if (!axera_models_[2]->SetInputTensorData("asr", asr.data(), asr.size()) || + !axera_models_[2]->SetInputTensorData("F0_pred", f0_pred.data(), + f0_pred.size()) || + !axera_models_[2]->SetInputTensorData("N_pred", n_pred.data(), + n_pred.size()) || + !axera_models_[2]->SetInputTensorData("ref_s", ref_s.data(), + ref_s.size()) || + !axera_models_[2]->SetInputTensorData("har", har.data(), har.size())) { return false; } - if (!model3_->Run()) { + if (!axera_models_[2]->Run()) { return false; } - std::vector x_out = model3_->GetOutputTensorData("x"); + std::vector x_out = axera_models_[2]->GetOutputTensorData("x"); PostprocessXToAudio(x_out, 23041, audio); if (is_doubled) { @@ -723,9 +775,7 @@ class OfflineTtsKokoroModelAxera::Impl { Ort::SessionOptions sess_opts_; Ort::AllocatorWithDefaultOptions allocator_; - std::unique_ptr model1_; - std::unique_ptr model2_; - std::unique_ptr model3_; + std::vector> axera_models_; std::unique_ptr model4_; std::vector model4_input_names_; diff --git a/sherpa-onnx/csrc/offline-tts-kokoro-model-config.cc b/sherpa-onnx/csrc/offline-tts-kokoro-model-config.cc index 06ab74d2b8..a523733381 100644 --- a/sherpa-onnx/csrc/offline-tts-kokoro-model-config.cc +++ b/sherpa-onnx/csrc/offline-tts-kokoro-model-config.cc @@ -45,9 +45,15 @@ bool OfflineTtsKokoroModelConfig::Validate() const { return false; } - if (!FileExists(model)) { - SHERPA_ONNX_LOGE("--kokoro-model: '%s' does not exist", model.c_str()); - return false; + if (!model.empty()) { + std::vector files; + SplitStringToVector(model, ",", false, &files); + for (const auto &f : files) { + if (!FileExists(f)) { + SHERPA_ONNX_LOGE("--kokoro-model: '%s' does not exist", model.c_str()); + return false; + } + } } if (tokens.empty()) { From e6fb94b10bcf170797816deb9806d354c9871048 Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Wed, 1 Apr 2026 16:12:06 +0800 Subject: [PATCH 06/10] Refactor input preparation to support variable repeat counts in OfflineTtsKokoroModelAxera --- .../axera/offline-tts-kokoro-model-axera.cc | 46 +++++++++++-------- 1 file changed, 26 insertions(+), 20 deletions(-) diff --git a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc index a0fe960d09..475db76fb0 100644 --- a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc +++ b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc @@ -470,10 +470,10 @@ class OfflineTtsKokoroModelAxera::Impl { float speed, std::vector &audio, int32_t &actual_content_frames, int32_t &total_frames) { - bool is_doubled = false; + int32_t repeat_count = 1; int32_t original_actual_len = actual_len; - PrepareInputIds(input_ids, actual_len, is_doubled); + PrepareInputIds(input_ids, actual_len, repeat_count); std::vector input_ids_i32(input_ids.begin(), input_ids.end()); @@ -577,11 +577,11 @@ class OfflineTtsKokoroModelAxera::Impl { std::vector x_out = axera_models_[2]->GetOutputTensorData("x"); PostprocessXToAudio(x_out, 23041, audio); - if (is_doubled) { + if (repeat_count > 1) { actual_content_frames = std::accumulate( pred_dur.begin(), pred_dur.begin() + original_actual_len, 0); - audio.erase(audio.begin() + audio.size() / 2, audio.end()); - total_frames = total_frames / 2; + audio.erase(audio.begin() + audio.size() / repeat_count, audio.end()); + total_frames = total_frames / repeat_count; } else { actual_content_frames = std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); @@ -730,31 +730,37 @@ class OfflineTtsKokoroModelAxera::Impl { } void PrepareInputIds(std::vector &input_ids, int32_t &actual_len, - bool &is_doubled) { - is_doubled = false; + int32_t &repeat_count) { + repeat_count = 1; int32_t original_actual_len = actual_len; - if (actual_len <= 32) { // DOUBLE_INPUT_THRESHOLD - is_doubled = true; + if (actual_len <= 32) { // REPEAT_INPUT_THRESHOLD + repeat_count = std::min(3, max_seq_len_ / std::max(actual_len, 1)); + if (repeat_count <= 1) { + return; + } + std::vector valid_content(input_ids.begin(), input_ids.begin() + actual_len); - std::vector input_ids_doubled; - input_ids_doubled.insert(input_ids_doubled.end(), valid_content.begin(), - valid_content.end()); - input_ids_doubled.insert(input_ids_doubled.end(), valid_content.begin(), - valid_content.end()); + std::vector repeated_input_ids; + repeated_input_ids.reserve(max_seq_len_); + + for (int32_t i = 0; i < repeat_count; ++i) { + repeated_input_ids.insert(repeated_input_ids.end(), + valid_content.begin(), valid_content.end()); + } - int32_t padding_len = max_seq_len_ - 2 * actual_len; + int32_t padding_len = max_seq_len_ - repeat_count * actual_len; if (padding_len > 0) { std::vector padding(padding_len, 0); - input_ids_doubled.insert(input_ids_doubled.end(), padding.begin(), - padding.end()); + repeated_input_ids.insert(repeated_input_ids.end(), padding.begin(), + padding.end()); } else { - input_ids_doubled.resize(max_seq_len_); + repeated_input_ids.resize(max_seq_len_); } - input_ids = input_ids_doubled; - actual_len = std::min(original_actual_len * 2, max_seq_len_); + input_ids = repeated_input_ids; + actual_len = std::min(original_actual_len * repeat_count, max_seq_len_); } } From f859ba9e9343bd45563ed103bb74db0a35862d6c Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Fri, 3 Apr 2026 17:28:29 +0800 Subject: [PATCH 07/10] Implement feature X to enhance user experience and optimize performance --- .../axera/offline-tts-kokoro-model-axera.cc | 1018 +++++++++-------- 1 file changed, 545 insertions(+), 473 deletions(-) diff --git a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc index 475db76fb0..928f72e7c3 100644 --- a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc +++ b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc @@ -6,12 +6,14 @@ #include #include +#include #include -#include +#include #include #include #include #include +#include #include #include #include @@ -28,7 +30,6 @@ #include "ax_engine_api.h" // NOLINT #include "ax_sys_api.h" // NOLINT -#include "kaldi-native-fbank/csrc/istft.h" #include "sherpa-onnx/csrc/axera/ax-engine-guard.h" #include "sherpa-onnx/csrc/axera/utils.h" #include "sherpa-onnx/csrc/file-utils.h" @@ -39,43 +40,182 @@ namespace sherpa_onnx { -static std::vector Sigmoid(const std::vector &x) { - std::vector result(x.size()); - for (size_t i = 0; i < x.size(); i++) { - result[i] = 1.0f / (1.0f + std::exp(-x[i])); +namespace { + +constexpr int32_t kKokoroSampleRate = 24000; +constexpr int32_t kKokoroVersion = 2; +constexpr int32_t kKokoroVoiceStyleLength = 510; +constexpr int32_t kKokoroStyleRank = 3; +constexpr int32_t kKokoroStyleMiddleDim = 1; +constexpr int32_t kKokoroTimbreDim = 128; +constexpr int32_t kHarmonicCount = 8; +constexpr int32_t kHarmonicDim = kHarmonicCount + 1; +constexpr int32_t kStftNfft = 20; +constexpr int32_t kStftHop = 5; +constexpr float kPi = 3.14159265358979323846f; +constexpr float kTwoPi = 2.0f * kPi; +constexpr float kSineAmp = 0.1f; +constexpr float kNoiseStd = 0.003f; +constexpr float kUnvoicedNoiseScale = kSineAmp / 3.0f; +constexpr float kVoicedThreshold = 10.0f; +constexpr char kDefaultVoice[] = "en-us"; +constexpr std::array kHarmonicMergeWeights = { + -0.08154187f, -0.18519667f, -0.18263398f, -0.17837206f, -0.09873895f, + 0.08264039f, 0.08743999f, -0.39068547f, -0.54774433f, +}; +constexpr float kHarmonicMergeBias = -0.02945026f; + +static int32_t GetNumElements(const std::vector &shape) { + if (shape.empty()) { + return 0; + } + + int32_t ans = 1; + for (auto dim : shape) { + ans *= dim; + } + + return ans; +} + +static std::string Basename(const std::string &path) { + auto pos = path.find_last_of("/\\"); + if (pos == std::string::npos) { + return path; } - return result; + + return path.substr(pos + 1); +} + +static int32_t ExtractIntegerAfterTag(const std::string &s, + const std::string &tag) { + auto pos = s.find(tag); + if (pos == std::string::npos) { + return -1; + } + + pos += tag.size(); + int32_t value = 0; + bool found = false; + while (pos < s.size() && std::isdigit(static_cast(s[pos]))) { + found = true; + value = value * 10 + (s[pos] - '0'); + ++pos; + } + + return found ? value : -1; +} + +static std::vector CropLastDim(const std::vector &src, + int32_t outer_dim, + int32_t src_last_dim, + int32_t dst_last_dim) { + std::vector ans(outer_dim * dst_last_dim); + for (int32_t i = 0; i != outer_dim; ++i) { + std::memcpy(ans.data() + i * dst_last_dim, + src.data() + i * src_last_dim, + dst_last_dim * sizeof(float)); + } + + return ans; } -template -static std::vector Argsort(const std::vector &v, int32_t len, - bool reverse) { - std::vector idx(len); - std::iota(idx.begin(), idx.end(), 0); - - if (!reverse) { - std::stable_sort(idx.begin(), idx.end(), - [&v](size_t i1, size_t i2) { return v[i1] < v[i2]; }); - } else { - std::stable_sort(idx.begin(), idx.end(), - [&v](size_t i1, size_t i2) { return v[i1] > v[i2]; }); +static std::vector PadLastDim(const std::vector &src, + int32_t outer_dim, int32_t src_last_dim, + int32_t dst_last_dim) { + std::vector ans(outer_dim * dst_last_dim, 0.0f); + for (int32_t i = 0; i != outer_dim; ++i) { + std::memcpy(ans.data() + i * dst_last_dim, + src.data() + i * src_last_dim, + src_last_dim * sizeof(float)); } - return idx; + return ans; } -template -static std::vector NpRepeat(const std::vector &v, - const std::vector ×) { - std::vector result; - for (size_t i = 0; i < times.size(); i++) { - for (int32_t n = 0; n < times[i]; n++) { - result.push_back(v[i]); +static std::vector PadAlignment(const std::vector &src, + int32_t src_token_len, + int32_t src_frame_len, + int32_t dst_token_len, + int32_t dst_frame_len) { + std::vector ans(dst_token_len * dst_frame_len, 0.0f); + for (int32_t i = 0; i != src_token_len; ++i) { + std::memcpy(ans.data() + i * dst_frame_len, + src.data() + i * src_frame_len, + src_frame_len * sizeof(float)); + } + + return ans; +} + +static std::vector BuildPaddedTextMask(int32_t token_len, + int32_t token_bucket) { + std::vector mask(token_bucket, 0); + for (int32_t i = token_len; i < token_bucket; ++i) { + mask[i] = 1; + } + + return mask; +} + +struct ModelPaths { + std::string encoder; + std::string duration_predictor; + std::string text_encoder; + std::string f0n_shared; + std::string f0n_head; + std::string decoder_front; + std::string vocoder; +}; + +static ModelPaths ParseModelPaths(const std::string &model_str) { + std::vector files; + SplitStringToVector(model_str, ",", false, &files); + + ModelPaths paths; + for (const auto &f : files) { + if (!FileExists(f)) { + SHERPA_ONNX_LOGE("Model file '%s' does not exist", f.c_str()); + SHERPA_ONNX_EXIT(-1); + } + + std::string name = Basename(f); + if (name.find("duration_predictor") != std::string::npos) { + paths.duration_predictor = f; + } else if (name.find("text_encoder_token_") != std::string::npos) { + paths.text_encoder = f; + } else if (name.find("encoder_token_") != std::string::npos) { + paths.encoder = f; + } else if (name.find("f0n_shared_frame_") != std::string::npos) { + paths.f0n_shared = f; + } else if (name.find("f0n_head_frame_") != std::string::npos) { + paths.f0n_head = f; + } else if (name.find("decoder_front") != std::string::npos) { + paths.decoder_front = f; + } else if (name.find("vocoder") != std::string::npos) { + paths.vocoder = f; } } - return result; + + if (paths.encoder.empty() || paths.duration_predictor.empty() || + paths.text_encoder.empty() || paths.f0n_shared.empty() || + paths.f0n_head.empty() || paths.decoder_front.empty() || + paths.vocoder.empty()) { + SHERPA_ONNX_LOGE( + "For the new Axera Kokoro pipeline, --kokoro-model must contain these " + "7 files in any order: encoder_token_*.axmodel, " + "duration_predictor.onnx, text_encoder_token_*_frame_*.axmodel, " + "f0n_shared_frame_*.axmodel, f0n_head_frame_*.axmodel, " + "decoder_front*.axmodel, vocoder*.axmodel. Given: %s", + model_str.c_str()); + SHERPA_ONNX_EXIT(-1); + } + + return paths; } +} // namespace + class AxeraModel { public: explicit AxeraModel(const std::string &filename) { @@ -123,7 +263,6 @@ class AxeraModel { int32_t i = it->second; size_t expected_size = io_info_->pInputs[i].nSize; size_t given_size = n * sizeof(T); - if (expected_size != given_size) { SHERPA_ONNX_LOGE( "Input tensor '%s' size mismatch. Expected %zu bytes, got %zu bytes", @@ -151,12 +290,11 @@ class AxeraModel { size_t out_elems = out_meta.nSize / sizeof(float); std::vector out(out_elems); std::memcpy(out.data(), out_buf.pVirAddr, out_meta.nSize); - return out; } bool Run() { - auto ret = AX_ENGINE_RunSync(handle_, &io_data_); + int ret = AX_ENGINE_RunSync(handle_, &io_data_); if (ret != 0) { SHERPA_ONNX_LOGE("AX_ENGINE_RunSync failed, ret = %d", ret); return false; @@ -196,23 +334,10 @@ class AxeraModel { std::unordered_map input_name_to_index_; std::unordered_map output_name_to_index_; - std::vector> input_tensor_shapes_; std::vector> output_tensor_shapes_; }; -static std::string GetModelPath(const std::string &model, const std::string &p, - bool is_onnx = false) { - std::string suffix = is_onnx ? ".onnx" : ".axmodel"; - - std::string path1 = model + "/kokoro_" + p + suffix; - if (FileExists(path1)) { - return path1; - } - - return model + "_" + p + suffix; -} - class OfflineTtsKokoroModelAxera::Impl { public: explicit Impl(const OfflineTtsModelConfig &config) @@ -220,7 +345,7 @@ class OfflineTtsKokoroModelAxera::Impl { env_(ORT_LOGGING_LEVEL_ERROR), sess_opts_(GetSessionOptions(config)), allocator_{} { - Init(config); + Init(config, [](const std::string &filename) { return ReadFile(filename); }); } template @@ -229,14 +354,14 @@ class OfflineTtsKokoroModelAxera::Impl { env_(ORT_LOGGING_LEVEL_ERROR), sess_opts_(GetSessionOptions(config)), allocator_{} { - Init(mgr, config); + Init(config, [mgr](const std::string &filename) { return ReadFile(mgr, filename); }); } const OfflineTtsKokoroModelMetaData &GetMetaData() const { return meta_data_; } - std::vector Run(std::vector input_ids, int64_t sid, + std::vector Run(const std::vector &input_ids, int64_t sid, float speed) { std::lock_guard lock(mutex_); @@ -245,19 +370,20 @@ class OfflineTtsKokoroModelAxera::Impl { } if (config_.kokoro.length_scale != 1 && speed == 1) { - speed = 1. / config_.kokoro.length_scale; + speed = 1.0f / config_.kokoro.length_scale; } int32_t phoneme_len = static_cast(input_ids.size()) - 2; - if (phoneme_len < 0) phoneme_len = 0; + if (phoneme_len < 0) { + phoneme_len = 0; + } std::vector ref_s = LoadVoiceEmbedding(static_cast(sid), phoneme_len); std::vector audio; - bool success = RunModels(input_ids, ref_s, speed, audio); - if (!success) { - SHERPA_ONNX_LOGE("Run models failed!"); + if (!RunModels(input_ids, ref_s, speed, &audio)) { + SHERPA_ONNX_LOGE("Run Kokoro Axera static pipeline failed"); return {}; } @@ -265,515 +391,447 @@ class OfflineTtsKokoroModelAxera::Impl { } private: - void Init(const OfflineTtsModelConfig &config) { - std::string model_str = config.kokoro.model; - - std::vector model_files; - SplitStringToVector(model_str, ",", false, &model_files); - - std::vector model_paths; - - if (config_.provider == "axera") { - // axera provider requires exactly 4 model files - if (model_files.size() != 4) { - SHERPA_ONNX_LOGE( - "For axera provider, model should contain exactly 4 files " - "separated by comma. Given %d files: %s", - static_cast(model_files.size()), model_str.c_str()); - SHERPA_ONNX_EXIT(-1); - } - - // Validate all 4 files exist - for (int32_t i = 0; i < 4; ++i) { - if (!FileExists(model_files[i])) { - SHERPA_ONNX_LOGE("Model file '%s' does not exist", - model_files[i].c_str()); - SHERPA_ONNX_EXIT(-1); - } - } - model_paths = model_files; - } else { + template + void Init(const OfflineTtsModelConfig &config, ReadBytes &&read_bytes) { + if (config_.provider != "axera") { SHERPA_ONNX_LOGE( "This model only supports axera provider. Please use provider=axera"); SHERPA_ONNX_EXIT(-1); } - // Load axera models (indices 0, 1, 2) - for (int32_t i = 0; i < 3; ++i) { - axera_models_.push_back(std::make_unique(model_paths[i])); - } + auto paths = ParseModelPaths(config.kokoro.model); + + auto encoder_buf = read_bytes(paths.encoder); + encoder_model_ = std::make_unique(encoder_buf.data(), + encoder_buf.size()); + + auto text_encoder_buf = read_bytes(paths.text_encoder); + text_encoder_model_ = std::make_unique( + text_encoder_buf.data(), text_encoder_buf.size()); + + auto f0n_shared_buf = read_bytes(paths.f0n_shared); + f0n_shared_model_ = + std::make_unique(f0n_shared_buf.data(), f0n_shared_buf.size()); + + auto f0n_head_buf = read_bytes(paths.f0n_head); + f0n_head_model_ = + std::make_unique(f0n_head_buf.data(), f0n_head_buf.size()); - // Load onnx model (index 3) - auto buf4 = ReadFile(model_paths[3]); - model4_ = std::make_unique(env_, buf4.data(), buf4.size(), - sess_opts_); - GetInputNames(model4_.get(), &model4_input_names_, - &model4_input_names_ptr_); - GetOutputNames(model4_.get(), &model4_output_names_, - &model4_output_names_ptr_); + auto decoder_front_buf = read_bytes(paths.decoder_front); + decoder_front_model_ = std::make_unique(decoder_front_buf.data(), + decoder_front_buf.size()); - auto voices_buf = ReadFile(config.kokoro.voices); + auto vocoder_buf = read_bytes(paths.vocoder); + vocoder_model_ = + std::make_unique(vocoder_buf.data(), vocoder_buf.size()); + + auto duration_buf = read_bytes(paths.duration_predictor); + duration_predictor_ = std::make_unique( + env_, duration_buf.data(), duration_buf.size(), sess_opts_); + GetInputNames(duration_predictor_.get(), &duration_input_names_, + &duration_input_names_ptr_); + GetOutputNames(duration_predictor_.get(), &duration_output_names_, + &duration_output_names_ptr_); + + InitModelShapes(paths); + + auto voices_buf = read_bytes(config.kokoro.voices); LoadVoices(voices_buf.data(), voices_buf.size()); } - template - void Init(Manager *mgr, const OfflineTtsModelConfig &config) { - std::string model_str = config.kokoro.model; - - std::vector model_files; - SplitStringToVector(model_str, ",", false, &model_files); - - std::vector model_paths; - - if (config_.provider == "axera") { - // axera provider requires exactly 4 model files - if (model_files.size() != 4) { - SHERPA_ONNX_LOGE( - "For axera provider, model should contain exactly 4 files " - "separated by comma. Given %d files: %s", - static_cast(model_files.size()), model_str.c_str()); - SHERPA_ONNX_EXIT(-1); - } - model_paths = model_files; - } else { - SHERPA_ONNX_LOGE( - "This model only supports axera provider. Please use provider=axera"); + void InitModelShapes(const ModelPaths &paths) { + auto encoder_input_shape = encoder_model_->TensorShape("input_ids"); + auto encoder_mask_shape = encoder_model_->TensorShape("text_mask"); + auto text_encoder_input_shape = text_encoder_model_->TensorShape("input_ids"); + auto text_encoder_aln_shape = text_encoder_model_->TensorShape("pred_aln_trg"); + auto f0n_shared_input_shape = f0n_shared_model_->TensorShape("en"); + auto f0n_head_ref_shape = f0n_head_model_->TensorShape("ref_s"); + auto f0n_head_output_shape = f0n_head_model_->TensorShape("F0_pred"); + auto decoder_asr_shape = decoder_front_model_->TensorShape("asr"); + auto decoder_timbre_shape = decoder_front_model_->TensorShape("timbre"); + auto vocoder_har_shape = vocoder_model_->TensorShape("har"); + auto vocoder_output_shape = vocoder_model_->TensorShape("waveform"); + + if (encoder_input_shape.size() != 2 || encoder_mask_shape.size() != 2 || + text_encoder_input_shape.size() != 2 || text_encoder_aln_shape.size() != 3 || + f0n_shared_input_shape.size() != 3 || f0n_head_ref_shape.size() != 2 || + f0n_head_output_shape.size() != 2 || decoder_asr_shape.size() != 3 || + decoder_timbre_shape.size() != 2 || vocoder_har_shape.size() != 3) { + SHERPA_ONNX_LOGE("Unexpected Kokoro Axera model shapes in: %s", + config_.kokoro.model.c_str()); SHERPA_ONNX_EXIT(-1); } - // Load axera models (indices 0, 1, 2) - for (int32_t i = 0; i < 3; ++i) { - auto buf = ReadFile(mgr, model_paths[i]); - axera_models_.push_back( - std::make_unique(buf.data(), buf.size())); + token_bucket_ = encoder_input_shape[1]; + frame_bucket_ = text_encoder_aln_shape[2]; + pitch_bucket_ = f0n_head_output_shape[1]; + ref_s_dim_ = f0n_head_ref_shape[1]; + timbre_dim_ = decoder_timbre_shape[1]; + har_channels_ = vocoder_har_shape[1]; + har_frames_ = vocoder_har_shape[2]; + sample_bucket_ = GetNumElements(vocoder_output_shape); + + if (token_bucket_ != text_encoder_input_shape[1] || + token_bucket_ != text_encoder_aln_shape[1] || + frame_bucket_ != f0n_shared_input_shape[2] || + frame_bucket_ != decoder_asr_shape[2]) { + SHERPA_ONNX_LOGE("Kokoro Axera bucket mismatch detected in: %s", + config_.kokoro.model.c_str()); + SHERPA_ONNX_EXIT(-1); } - // Load onnx model (index 3) - auto buf4 = ReadFile(mgr, model_paths[3]); - model4_ = std::make_unique(env_, buf4.data(), buf4.size(), - sess_opts_); - GetInputNames(model4_.get(), &model4_input_names_, - &model4_input_names_ptr_); - GetOutputNames(model4_.get(), &model4_output_names_, - &model4_output_names_ptr_); - - auto voices_buf = ReadFile(mgr, config.kokoro.voices); - LoadVoices(voices_buf.data(), voices_buf.size()); - } + if (ref_s_dim_ <= 0 || timbre_dim_ <= 0 || timbre_dim_ > ref_s_dim_) { + SHERPA_ONNX_LOGE("Invalid Kokoro style dims. ref_s=%d, timbre=%d", + ref_s_dim_, timbre_dim_); + SHERPA_ONNX_EXIT(-1); + } - void LoadVoices(const char *voices_data, size_t voices_data_length) { - // Get metadata from model4 (ONNX model) - Ort::ModelMetadata meta_data = model4_->GetModelMetadata(); - Ort::AllocatorWithDefaultOptions allocator; - - SHERPA_ONNX_READ_META_DATA(meta_data_.sample_rate, "sample_rate"); - SHERPA_ONNX_READ_META_DATA_WITH_DEFAULT(meta_data_.version, "version", 1); - SHERPA_ONNX_READ_META_DATA(meta_data_.num_speakers, "n_speakers"); - SHERPA_ONNX_READ_META_DATA(meta_data_.has_espeak, "has_espeak"); - SHERPA_ONNX_READ_META_DATA_STR_WITH_DEFAULT(meta_data_.voice, "voice", - "en-us"); - SHERPA_ONNX_READ_META_DATA_VEC(style_dim_, "style_dim"); - - if (style_dim_.size() != 3) { - SHERPA_ONNX_LOGE("style_dim should be 3-d, given: %d", - static_cast(style_dim_.size())); + if (har_channels_ != kStftNfft + 2) { + SHERPA_ONNX_LOGE("Expected har channels %d. Given: %d", kStftNfft + 2, + har_channels_); SHERPA_ONNX_EXIT(-1); } - if (style_dim_[1] != 1) { - SHERPA_ONNX_LOGE("style_dim[1] should be 1, given: %d", style_dim_[1]); + if (sample_bucket_ % pitch_bucket_ != 0 || sample_bucket_ % frame_bucket_ != 0) { + SHERPA_ONNX_LOGE( + "Unexpected Kokoro Axera bucket relation. sample=%d pitch=%d frame=%d", + sample_bucket_, pitch_bucket_, frame_bucket_); SHERPA_ONNX_EXIT(-1); } - meta_data_.max_token_len = style_dim_[0]; + upsample_scale_ = sample_bucket_ / pitch_bucket_; + samples_per_frame_ = sample_bucket_ / frame_bucket_; + int32_t expected_har_frames = sample_bucket_ / kStftHop + 1; + if (har_frames_ != expected_har_frames) { + SHERPA_ONNX_LOGE("Expected har frames %d. Given: %d", expected_har_frames, + har_frames_); + SHERPA_ONNX_EXIT(-1); + } - int32_t actual_num_floats = voices_data_length / sizeof(float); - int32_t expected_num_floats = - style_dim_[0] * style_dim_[2] * meta_data_.num_speakers; + if (ExtractIntegerAfterTag(Basename(paths.encoder), "encoder_token_") != -1 && + ExtractIntegerAfterTag(Basename(paths.encoder), "encoder_token_") != + token_bucket_) { + SHERPA_ONNX_LOGE("Encoder filename bucket does not match model shape: %s", + paths.encoder.c_str()); + SHERPA_ONNX_EXIT(-1); + } + } - if (actual_num_floats != expected_num_floats) { - SHERPA_ONNX_LOGE( - "Corrupted voices file. Expected #floats: %d, actual: %d", - expected_num_floats, actual_num_floats); + void LoadVoices(const char *voices_data, size_t voices_data_length) { + meta_data_.sample_rate = kKokoroSampleRate; + meta_data_.version = kKokoroVersion; + meta_data_.has_espeak = 1; + meta_data_.voice = config_.kokoro.lang.empty() ? kDefaultVoice : config_.kokoro.lang; + meta_data_.max_token_len = kKokoroVoiceStyleLength; + + style_dim_ = {kKokoroVoiceStyleLength, kKokoroStyleMiddleDim, ref_s_dim_}; + if (style_dim_.size() != kKokoroStyleRank || style_dim_[1] != 1) { + SHERPA_ONNX_LOGE("Unexpected Kokoro style_dim"); SHERPA_ONNX_EXIT(-1); } - styles_ = std::vector( - reinterpret_cast(voices_data), - reinterpret_cast(voices_data) + actual_num_floats); - auto shape1 = axera_models_[0]->TensorShape("input_ids"); - if (shape1.size() < 2) { - SHERPA_ONNX_LOGE("Unexpected shape rank for input_ids: %d", - static_cast(shape1.size())); + int32_t actual_num_floats = static_cast(voices_data_length / sizeof(float)); + int32_t floats_per_speaker = style_dim_[0] * style_dim_[2]; + if (floats_per_speaker <= 0 || actual_num_floats % floats_per_speaker != 0) { + SHERPA_ONNX_LOGE( + "Corrupted voices file '%s'. Expected a multiple of %d floats. Got %d", + config_.kokoro.voices.c_str(), floats_per_speaker, actual_num_floats); SHERPA_ONNX_EXIT(-1); } - max_seq_len_ = shape1[1]; + meta_data_.num_speakers = actual_num_floats / floats_per_speaker; + styles_ = std::vector(reinterpret_cast(voices_data), + reinterpret_cast(voices_data) + + actual_num_floats); } - std::vector LoadVoiceEmbedding(int32_t sid, int32_t phoneme_len) { - int32_t style_dim0 = style_dim_[0]; - int32_t style_dim1 = style_dim_[2]; - phoneme_len = std::max(phoneme_len, 0); + std::vector LoadVoiceEmbedding(int32_t sid, int32_t phoneme_len) const { + int32_t style_len = style_dim_[0]; + int32_t style_dim = style_dim_[2]; sid = std::max(sid, 0); if (meta_data_.num_speakers > 0) { sid = std::min(sid, meta_data_.num_speakers - 1); } - std::vector ref_s(style_dim1); - if (phoneme_len < style_dim0) { - const float *p = styles_.data() + sid * style_dim0 * style_dim1 + - phoneme_len * style_dim1; - std::copy(p, p + style_dim1, ref_s.begin()); - } else { - int32_t idx = style_dim0 / 2; - const float *p = - styles_.data() + sid * style_dim0 * style_dim1 + idx * style_dim1; - std::copy(p, p + style_dim1, ref_s.begin()); - } + phoneme_len = std::max(phoneme_len, 0); + + std::vector ref_s(style_dim); + int32_t index = phoneme_len < style_len ? phoneme_len : style_len / 2; + const float *src = styles_.data() + sid * style_len * style_dim + + index * style_dim; + std::copy(src, src + style_dim, ref_s.begin()); return ref_s; } - bool RunModels(std::vector &input_ids, - const std::vector &ref_s, float speed, - std::vector &audio) { - int32_t actual_len = input_ids.size(); - if (actual_len > max_seq_len_) { - SHERPA_ONNX_LOGE( - "Input ids length %d exceeds max_seq_len %d, truncating.", actual_len, - max_seq_len_); - input_ids.resize(max_seq_len_); - actual_len = max_seq_len_; - } - - int32_t padding_len = max_seq_len_ - actual_len; - if (padding_len > 0) { - std::vector padding(padding_len, 0); - input_ids.insert(input_ids.end(), padding.begin(), padding.end()); + std::vector BuildHar(const std::vector &f0_pred) const { + std::vector phase_coarse(pitch_bucket_ * kHarmonicDim, 0.0f); + std::array cumulative{}; + + for (int32_t t = 0; t != pitch_bucket_; ++t) { + float base_f0 = f0_pred[t]; + for (int32_t h = 0; h != kHarmonicDim; ++h) { + float harmonic = base_f0 * static_cast(h + 1); + float rad = std::fmod(harmonic / static_cast(kKokoroSampleRate), 1.0f); + if (rad < 0) { + rad += 1.0f; + } + cumulative[h] += rad; + phase_coarse[t * kHarmonicDim + h] = + cumulative[h] * kTwoPi * static_cast(upsample_scale_); + } } - int32_t actual_content_frames = 0; - int32_t total_frames = 0; - if (!InferenceSingleChunk(input_ids, ref_s, actual_len, speed, audio, - actual_content_frames, total_frames)) { - return false; - } + std::vector sine_merge(sample_bucket_); + std::mt19937 rng(0); + std::normal_distribution normal(0.0f, 1.0f); - TrimAudioByContent(audio, actual_content_frames, total_frames, actual_len); + for (int32_t sample_index = 0; sample_index != sample_bucket_; ++sample_index) { + int32_t coarse_index = sample_index / upsample_scale_; + float coarse_f0 = f0_pred[coarse_index]; + float uv = coarse_f0 > kVoicedThreshold ? 1.0f : 0.0f; + float noise_amp = uv * kNoiseStd + (1.0f - uv) * kUnvoicedNoiseScale; - return true; - } + float src = (static_cast(sample_index) + 0.5f) / + static_cast(upsample_scale_) - + 0.5f; + if (src < 0.0f) { + src = 0.0f; + } - bool InferenceSingleChunk(std::vector &input_ids, - const std::vector &ref_s, int32_t actual_len, - float speed, std::vector &audio, - int32_t &actual_content_frames, - int32_t &total_frames) { - int32_t repeat_count = 1; - int32_t original_actual_len = actual_len; + int32_t left = static_cast(std::floor(src)); + int32_t right = std::min(left + 1, pitch_bucket_ - 1); + float w = src - static_cast(left); + if (left >= pitch_bucket_ - 1) { + left = pitch_bucket_ - 1; + right = left; + w = 0.0f; + } - PrepareInputIds(input_ids, actual_len, repeat_count); + float merged = kHarmonicMergeBias; + for (int32_t h = 0; h != kHarmonicDim; ++h) { + float phase_left = phase_coarse[left * kHarmonicDim + h]; + float phase_right = phase_coarse[right * kHarmonicDim + h]; + float phase = phase_left * (1.0f - w) + phase_right * w; + float sine = std::sin(phase) * kSineAmp; + float noise = noise_amp * normal(rng); + float sine_wave = sine * uv + noise; + merged += sine_wave * kHarmonicMergeWeights[h]; + } - std::vector input_ids_i32(input_ids.begin(), input_ids.end()); + sine_merge[sample_index] = std::tanh(merged); + } - std::vector text_mask; - ComputeExternalPreprocessing(actual_len, text_mask); + std::vector padded(sample_bucket_ + kStftNfft, 0.0f); + int32_t pad = kStftNfft / 2; + std::copy(sine_merge.begin(), sine_merge.end(), padded.begin() + pad); + std::fill(padded.begin(), padded.begin() + pad, sine_merge.front()); + std::fill(padded.end() - pad, padded.end(), sine_merge.back()); - std::vector text_mask_u8(text_mask.begin(), text_mask.end()); + std::array window{}; + std::array, kStftNfft / 2 + 1> cos_table{}; + std::array, kStftNfft / 2 + 1> sin_table{}; - if (!axera_models_[0]->SetInputTensorData("input_ids", input_ids_i32.data(), - input_ids_i32.size()) || - !axera_models_[0]->SetInputTensorData("ref_s", ref_s.data(), - ref_s.size()) || - !axera_models_[0]->SetInputTensorData("text_mask", text_mask_u8.data(), - text_mask_u8.size())) { - return false; + for (int32_t n = 0; n != kStftNfft; ++n) { + window[n] = 0.5f - 0.5f * std::cos(kTwoPi * n / static_cast(kStftNfft)); } - if (!axera_models_[0]->Run()) { - return false; + for (int32_t k = 0; k != kStftNfft / 2 + 1; ++k) { + for (int32_t n = 0; n != kStftNfft; ++n) { + float angle = kTwoPi * k * n / static_cast(kStftNfft); + cos_table[k][n] = std::cos(angle); + sin_table[k][n] = -std::sin(angle); + } } - std::vector duration = - axera_models_[0]->GetOutputTensorData("duration"); - std::vector d = axera_models_[0]->GetOutputTensorData("d"); - - std::vector pred_dur; - ProcessDuration(duration, actual_len, speed, pred_dur, total_frames); - - std::vector pred_aln_trg = - CreateAlignmentMatrix(pred_dur, total_frames); - - int32_t d_shape_2 = 640; - std::vector en(total_frames * d_shape_2, 0.0f); + std::vector har(har_channels_ * har_frames_); + for (int32_t frame = 0; frame != har_frames_; ++frame) { + const float *segment = padded.data() + frame * kStftHop; + for (int32_t k = 0; k != kStftNfft / 2 + 1; ++k) { + float real = 0.0f; + float imag = 0.0f; + for (int32_t n = 0; n != kStftNfft; ++n) { + float value = segment[n] * window[n]; + real += value * cos_table[k][n]; + imag += value * sin_table[k][n]; + } - for (int32_t j = 0; j < d_shape_2; ++j) { - for (int32_t k = 0; k < total_frames; ++k) { - float sum = 0.0f; - for (int32_t i = 0; i < max_seq_len_; ++i) { - sum += d[i * d_shape_2 + j] * pred_aln_trg[i * total_frames + k]; + har[k * har_frames_ + frame] = + std::sqrt(real * real + imag * imag + 1e-14f); + float phase = std::atan2(imag, real); + if (imag == 0.0f && real < 0.0f) { + phase = kPi; } - en[j * total_frames + k] = sum; + har[(k + kStftNfft / 2 + 1) * har_frames_ + frame] = phase; } } - std::vector text_mask_float(text_mask.begin(), text_mask.end()); + return har; + } - if (!axera_models_[1]->SetInputTensorData("en", en.data(), en.size()) || - !axera_models_[1]->SetInputTensorData("ref_s", ref_s.data(), - ref_s.size()) || - !axera_models_[1]->SetInputTensorData("input_ids", input_ids_i32.data(), - input_ids_i32.size()) || - !axera_models_[1]->SetInputTensorData( - "text_mask", text_mask_float.data(), text_mask_float.size()) || - !axera_models_[1]->SetInputTensorData( - "pred_aln_trg", pred_aln_trg.data(), pred_aln_trg.size())) { + bool RunModels(const std::vector &input_ids, + const std::vector &ref_s, float speed, + std::vector *audio) { + int32_t token_length = static_cast(input_ids.size()); + if (token_length > token_bucket_) { + SHERPA_ONNX_LOGE( + "Token length %d exceeds Axera token bucket %d. Use a shorter text " + "or export a larger static bucket.", + token_length, token_bucket_); return false; } - if (!axera_models_[1]->Run()) { + std::vector padded_input_ids(token_bucket_, 0); + for (int32_t i = 0; i != token_length; ++i) { + padded_input_ids[i] = static_cast(input_ids[i]); + } + + std::vector padded_text_mask = + BuildPaddedTextMask(token_length, token_bucket_); + + if (!encoder_model_->SetInputTensorData("input_ids", padded_input_ids.data(), + padded_input_ids.size()) || + !encoder_model_->SetInputTensorData("text_mask", padded_text_mask.data(), + padded_text_mask.size()) || + !encoder_model_->Run()) { return false; } - std::vector f0_pred = - axera_models_[1]->GetOutputTensorData("F0_pred"); - std::vector n_pred = axera_models_[1]->GetOutputTensorData("N_pred"); - std::vector asr = axera_models_[1]->GetOutputTensorData("asr"); + std::vector padded_d_en = encoder_model_->GetOutputTensorData("d_en"); + std::vector d_en = + CropLastDim(padded_d_en, 512, token_bucket_, token_length); auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); - std::array f0_shape = {1, static_cast(f0_pred.size())}; - Ort::Value f0_tensor = - Ort::Value::CreateTensor(memory_info, f0_pred.data(), f0_pred.size(), - f0_shape.data(), f0_shape.size()); - - auto out4 = model4_->Run({}, model4_input_names_ptr_.data(), &f0_tensor, 1, - model4_output_names_ptr_.data(), - model4_output_names_ptr_.size()); - float *p_har = out4[0].GetTensorMutableData(); - auto har_shape = out4[0].GetTensorTypeAndShapeInfo().GetShape(); - int64_t har_num_elements = 1; - for (auto d : har_shape) { - har_num_elements *= d; - } - std::vector har(p_har, p_har + har_num_elements); - - if (!axera_models_[2]->SetInputTensorData("asr", asr.data(), asr.size()) || - !axera_models_[2]->SetInputTensorData("F0_pred", f0_pred.data(), - f0_pred.size()) || - !axera_models_[2]->SetInputTensorData("N_pred", n_pred.data(), - n_pred.size()) || - !axera_models_[2]->SetInputTensorData("ref_s", ref_s.data(), - ref_s.size()) || - !axera_models_[2]->SetInputTensorData("har", har.data(), har.size())) { + + std::array d_en_shape = {1, 512, token_length}; + std::array ref_s_shape = {1, ref_s_dim_}; + std::array input_lengths_shape = {1}; + std::array duration_text_mask_shape = {1, token_length}; + std::array speed_shape = {1}; + + std::array input_lengths = {token_length}; + std::unique_ptr duration_text_mask(new bool[token_length]()); + + Ort::Value d_en_tensor = Ort::Value::CreateTensor( + memory_info, d_en.data(), d_en.size(), d_en_shape.data(), d_en_shape.size()); + Ort::Value ref_s_tensor = Ort::Value::CreateTensor( + memory_info, const_cast(ref_s.data()), ref_s.size(), + ref_s_shape.data(), ref_s_shape.size()); + Ort::Value input_lengths_tensor = Ort::Value::CreateTensor( + memory_info, input_lengths.data(), input_lengths.size(), + input_lengths_shape.data(), input_lengths_shape.size()); + Ort::Value duration_text_mask_tensor = Ort::Value::CreateTensor( + memory_info, duration_text_mask.get(), token_length, + duration_text_mask_shape.data(), duration_text_mask_shape.size()); + Ort::Value speed_tensor = Ort::Value::CreateTensor( + memory_info, &speed, 1, speed_shape.data(), speed_shape.size()); + + std::array duration_inputs = { + std::move(d_en_tensor), std::move(ref_s_tensor), + std::move(input_lengths_tensor), std::move(duration_text_mask_tensor), + std::move(speed_tensor)}; + + auto duration_outputs = duration_predictor_->Run( + {}, duration_input_names_ptr_.data(), duration_inputs.data(), + duration_inputs.size(), duration_output_names_ptr_.data(), + duration_output_names_ptr_.size()); + + auto pred_aln_shape = duration_outputs[2].GetTensorTypeAndShapeInfo().GetShape(); + auto en_shape = duration_outputs[3].GetTensorTypeAndShapeInfo().GetShape(); + if (pred_aln_shape.size() != 3 || en_shape.size() != 3) { + SHERPA_ONNX_LOGE("Unexpected duration predictor output shapes"); return false; } - if (!axera_models_[2]->Run()) { + int32_t frame_length = static_cast(pred_aln_shape[2]); + if (frame_length > frame_bucket_) { + SHERPA_ONNX_LOGE( + "Frame length %d exceeds Axera frame bucket %d. Use a shorter text, " + "increase speed, or export a larger static bucket.", + frame_length, frame_bucket_); return false; } - std::vector x_out = axera_models_[2]->GetOutputTensorData("x"); - PostprocessXToAudio(x_out, 23041, audio); - - if (repeat_count > 1) { - actual_content_frames = std::accumulate( - pred_dur.begin(), pred_dur.begin() + original_actual_len, 0); - audio.erase(audio.begin() + audio.size() / repeat_count, audio.end()); - total_frames = total_frames / repeat_count; - } else { - actual_content_frames = - std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); - } - return true; - } - - void ProcessDuration(const std::vector &duration, int32_t actual_len, - float speed, std::vector &pred_dur, - int32_t &total_frames) { - std::vector pred_dur_original(actual_len, 0); - std::vector duration_processed = Sigmoid(duration); - - int32_t dur_shape_2 = 50; - for (int32_t i = 0; i < actual_len; i++) { - float sum = 0; - for (int32_t n = 0; n < dur_shape_2; n++) { - sum += duration_processed[i * dur_shape_2 + n]; - } - sum /= speed; - pred_dur_original[i] = - static_cast(std::max(1.f, std::round(sum))); - } - - std::vector pred_dur_padding(max_seq_len_ - actual_len, 0); - pred_dur = pred_dur_original; - pred_dur.insert(pred_dur.end(), pred_dur_padding.begin(), - pred_dur_padding.end()); - - int32_t fixed_total_frames = max_seq_len_ * 2; - int32_t actual_frames = - std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); - int32_t diff = fixed_total_frames - actual_frames; - - if (diff < 0) { - auto indices = Argsort(pred_dur, actual_len, true); - int32_t decreased = 0; - for (auto idx : indices) { - if (pred_dur[idx] > 1 && decreased < std::abs(diff)) { - pred_dur[idx]--; - decreased++; - } - if (decreased >= std::abs(diff)) break; - } + const float *pred_aln_ptr = duration_outputs[2].GetTensorData(); + int32_t pred_aln_elems = + static_cast(duration_outputs[2].GetTensorTypeAndShapeInfo().GetElementCount()); + std::vector pred_aln_trg(pred_aln_ptr, pred_aln_ptr + pred_aln_elems); + std::vector padded_alignment = PadAlignment( + pred_aln_trg, token_length, frame_length, token_bucket_, frame_bucket_); + + if (!text_encoder_model_->SetInputTensorData("input_ids", padded_input_ids.data(), + padded_input_ids.size()) || + !text_encoder_model_->SetInputTensorData("pred_aln_trg", + padded_alignment.data(), + padded_alignment.size()) || + !text_encoder_model_->SetInputTensorData("text_mask", + padded_text_mask.data(), + padded_text_mask.size()) || + !text_encoder_model_->Run()) { + return false; } - actual_frames = - std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); - int32_t remaining_frames = fixed_total_frames - actual_frames; - int32_t padding_len = max_seq_len_ - actual_len; - - if (remaining_frames > 0 && padding_len > 0) { - int32_t frames_per_padding = remaining_frames / padding_len; - int32_t remainder = remaining_frames % padding_len; + std::vector asr = text_encoder_model_->GetOutputTensorData("asr"); - for (int32_t i = actual_len; i < static_cast(pred_dur.size()); - i++) - pred_dur[i] = frames_per_padding; + const float *en_ptr = duration_outputs[3].GetTensorData(); + int32_t en_elems = static_cast( + duration_outputs[3].GetTensorTypeAndShapeInfo().GetElementCount()); + std::vector en(en_ptr, en_ptr + en_elems); + std::vector padded_en = PadLastDim(en, 640, frame_length, frame_bucket_); - if (remainder > 0) { - for (int32_t i = actual_len; i < actual_len + remainder; i++) - pred_dur[i] += 1; - } + if (!f0n_shared_model_->SetInputTensorData("en", padded_en.data(), + padded_en.size()) || + !f0n_shared_model_->Run()) { + return false; } - total_frames = std::accumulate(pred_dur.begin(), pred_dur.end(), 0); - } - - std::vector CreateAlignmentMatrix(const std::vector &pred_dur, - int32_t total_frames) { - std::vector seq_range(max_seq_len_); - std::iota(seq_range.begin(), seq_range.end(), 0); - auto indices = NpRepeat(seq_range, pred_dur); - - std::vector pred_aln_trg(max_seq_len_ * total_frames, 0.0f); - if (!indices.empty()) { - int32_t col = 0; - for (auto i : indices) { - pred_aln_trg[i * total_frames + col] = 1.0f; - col++; - } + std::vector shared = f0n_shared_model_->GetOutputTensorData("shared"); + if (!f0n_head_model_->SetInputTensorData("shared", shared.data(), shared.size()) || + !f0n_head_model_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()) || + !f0n_head_model_->Run()) { + return false; } - return pred_aln_trg; - } + std::vector f0_pred = f0n_head_model_->GetOutputTensorData("F0_pred"); + std::vector n_pred = f0n_head_model_->GetOutputTensorData("N_pred"); + std::vector timbre(ref_s.begin(), ref_s.begin() + timbre_dim_); - void PostprocessXToAudio(const std::vector &x, int32_t num_frames, - std::vector &audio) { - int32_t n_fft = 20; - int32_t hop_length = 5; - int32_t half_n_fft = n_fft / 2 + 1; - - std::vector spec_part(half_n_fft * num_frames); - std::vector phase_part(half_n_fft * num_frames); - std::vector cos_part(half_n_fft * num_frames); - spec_part.assign(x.begin(), x.begin() + half_n_fft * num_frames); - phase_part.assign(x.begin() + half_n_fft * num_frames, x.end()); - - for (int32_t i = 0; i < half_n_fft * num_frames; i++) { - spec_part[i] = std::exp(spec_part[i]); - phase_part[i] = std::sin(phase_part[i]); - cos_part[i] = std::sqrt( - 1.f - std::max(0.f, std::min(phase_part[i] * phase_part[i], 1.0f))); - } - - knf::StftResult stft_result; - stft_result.num_frames = num_frames; - stft_result.real.resize(half_n_fft * num_frames); - stft_result.imag.resize(half_n_fft * num_frames); - - for (int32_t i = 0; i < half_n_fft; i++) { - for (int32_t n = 0; n < num_frames; n++) { - float spec = spec_part[i * num_frames + n]; - float real_part = spec * cos_part[i * num_frames + n]; - float imag_part = spec * phase_part[i * num_frames + n]; - - // StftResult layout is [frame_index * half_n_fft + bin] - stft_result.real[n * half_n_fft + i] = real_part; - stft_result.imag[n * half_n_fft + i] = imag_part; - } + if (!decoder_front_model_->SetInputTensorData("asr", asr.data(), asr.size()) || + !decoder_front_model_->SetInputTensorData("F0_pred", f0_pred.data(), + f0_pred.size()) || + !decoder_front_model_->SetInputTensorData("N_pred", n_pred.data(), + n_pred.size()) || + !decoder_front_model_->SetInputTensorData("timbre", timbre.data(), + timbre.size()) || + !decoder_front_model_->Run()) { + return false; } - knf::StftConfig stft_config; - stft_config.n_fft = n_fft; - stft_config.hop_length = hop_length; - stft_config.win_length = n_fft; - stft_config.center = true; - stft_config.normalized = false; - stft_config.window_type = "hann"; - stft_config.pad_mode = "reflect"; - - knf::IStft istft(stft_config); - audio = istft.Compute(stft_result); - } + std::vector decoder_state = + decoder_front_model_->GetOutputTensorData("decoder_state"); + std::vector har = BuildHar(f0_pred); - void TrimAudioByContent(std::vector &audio, - int32_t actual_content_frames, int32_t total_frames, - int32_t actual_len) { - int32_t padding_len = max_seq_len_ - actual_len; - if (padding_len > 0) { - float content_ratio = actual_content_frames * 1.0f / total_frames; - int32_t audio_len_to_keep = - static_cast(audio.size() * content_ratio); - audio.resize(audio_len_to_keep); + if (!vocoder_model_->SetInputTensorData("decoder_state", decoder_state.data(), + decoder_state.size()) || + !vocoder_model_->SetInputTensorData("timbre", timbre.data(), + timbre.size()) || + !vocoder_model_->SetInputTensorData("har", har.data(), har.size()) || + !vocoder_model_->Run()) { + return false; } - } - - void PrepareInputIds(std::vector &input_ids, int32_t &actual_len, - int32_t &repeat_count) { - repeat_count = 1; - int32_t original_actual_len = actual_len; - - if (actual_len <= 32) { // REPEAT_INPUT_THRESHOLD - repeat_count = std::min(3, max_seq_len_ / std::max(actual_len, 1)); - if (repeat_count <= 1) { - return; - } - - std::vector valid_content(input_ids.begin(), - input_ids.begin() + actual_len); - std::vector repeated_input_ids; - repeated_input_ids.reserve(max_seq_len_); - for (int32_t i = 0; i < repeat_count; ++i) { - repeated_input_ids.insert(repeated_input_ids.end(), - valid_content.begin(), valid_content.end()); - } - - int32_t padding_len = max_seq_len_ - repeat_count * actual_len; - if (padding_len > 0) { - std::vector padding(padding_len, 0); - repeated_input_ids.insert(repeated_input_ids.end(), padding.begin(), - padding.end()); - } else { - repeated_input_ids.resize(max_seq_len_); - } - - input_ids = repeated_input_ids; - actual_len = std::min(original_actual_len * repeat_count, max_seq_len_); + *audio = vocoder_model_->GetOutputTensorData("waveform"); + int32_t sample_length = frame_length * samples_per_frame_; + if (sample_length < static_cast(audio->size())) { + audio->resize(sample_length); } - } - void ComputeExternalPreprocessing(int32_t actual_len, - std::vector &text_mask) { - text_mask.resize(max_seq_len_); - for (int32_t i = 0; i < max_seq_len_; i++) { - text_mask[i] = (i >= actual_len) ? 1 : 0; - } + return true; } private: - std::mutex mutex_; + mutable std::mutex mutex_; AxEngineGuard ax_engine_guard_; OfflineTtsModelConfig config_; @@ -781,19 +839,33 @@ class OfflineTtsKokoroModelAxera::Impl { Ort::SessionOptions sess_opts_; Ort::AllocatorWithDefaultOptions allocator_; - std::vector> axera_models_; + std::unique_ptr encoder_model_; + std::unique_ptr text_encoder_model_; + std::unique_ptr f0n_shared_model_; + std::unique_ptr f0n_head_model_; + std::unique_ptr decoder_front_model_; + std::unique_ptr vocoder_model_; - std::unique_ptr model4_; - std::vector model4_input_names_; - std::vector model4_input_names_ptr_; - std::vector model4_output_names_; - std::vector model4_output_names_ptr_; + std::unique_ptr duration_predictor_; + std::vector duration_input_names_; + std::vector duration_input_names_ptr_; + std::vector duration_output_names_; + std::vector duration_output_names_ptr_; OfflineTtsKokoroModelMetaData meta_data_; std::vector style_dim_; std::vector styles_; - int32_t max_seq_len_; + int32_t token_bucket_ = 0; + int32_t frame_bucket_ = 0; + int32_t pitch_bucket_ = 0; + int32_t sample_bucket_ = 0; + int32_t ref_s_dim_ = 0; + int32_t timbre_dim_ = 0; + int32_t har_channels_ = 0; + int32_t har_frames_ = 0; + int32_t upsample_scale_ = 0; + int32_t samples_per_frame_ = 0; }; OfflineTtsKokoroModelAxera::~OfflineTtsKokoroModelAxera() = default; From b4b4515e3ec5b9c68c313a02fd5a32c63c6eafbd Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Fri, 3 Apr 2026 17:43:38 +0800 Subject: [PATCH 08/10] Implement feature X to enhance user experience and optimize performance --- .../axcl/offline-tts-kokoro-model-axcl.cc | 1034 +++++++++-------- 1 file changed, 565 insertions(+), 469 deletions(-) diff --git a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc index f415e7a679..8afd50202b 100644 --- a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc +++ b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc @@ -6,11 +6,14 @@ #include #include +#include #include -#include +#include #include #include +#include #include +#include #include #include #include @@ -24,7 +27,6 @@ #include "rawfile/raw_file_manager.h" #endif -#include "kaldi-native-fbank/csrc/istft.h" #include "sherpa-onnx/csrc/axcl/axcl-model.h" #include "sherpa-onnx/csrc/file-utils.h" #include "sherpa-onnx/csrc/macros.h" @@ -34,69 +36,190 @@ namespace sherpa_onnx { -static std::vector Sigmoid(const std::vector &x) { - std::vector result(x.size()); - for (size_t i = 0; i < x.size(); i++) { - result[i] = 1.0f / (1.0f + std::exp(-x[i])); +namespace { + +constexpr int32_t kKokoroSampleRate = 24000; +constexpr int32_t kKokoroVersion = 2; +constexpr int32_t kKokoroVoiceStyleLength = 510; +constexpr int32_t kKokoroStyleRank = 3; +constexpr int32_t kKokoroStyleMiddleDim = 1; +constexpr int32_t kKokoroTimbreDim = 128; +constexpr int32_t kHarmonicCount = 8; +constexpr int32_t kHarmonicDim = kHarmonicCount + 1; +constexpr int32_t kStftNfft = 20; +constexpr int32_t kStftHop = 5; +constexpr float kPi = 3.14159265358979323846f; +constexpr float kTwoPi = 2.0f * kPi; +constexpr float kSineAmp = 0.1f; +constexpr float kNoiseStd = 0.003f; +constexpr float kUnvoicedNoiseScale = kSineAmp / 3.0f; +constexpr float kVoicedThreshold = 10.0f; +constexpr char kDefaultVoice[] = "en-us"; +constexpr std::array kHarmonicMergeWeights = { + -0.08154187f, -0.18519667f, -0.18263398f, -0.17837206f, -0.09873895f, + 0.08264039f, 0.08743999f, -0.39068547f, -0.54774433f, +}; +constexpr float kHarmonicMergeBias = -0.02945026f; + +static int32_t GetNumElements(const std::vector &shape) { + if (shape.empty()) { + return 0; } - return result; + + int32_t ans = 1; + for (auto dim : shape) { + ans *= dim; + } + + return ans; } -template -static std::vector Argsort(const std::vector &v, int32_t len, - bool reverse) { - std::vector idx(len); - std::iota(idx.begin(), idx.end(), 0); - - if (!reverse) { - std::stable_sort(idx.begin(), idx.end(), - [&v](size_t i1, size_t i2) { return v[i1] < v[i2]; }); - } else { - std::stable_sort(idx.begin(), idx.end(), - [&v](size_t i1, size_t i2) { return v[i1] > v[i2]; }); +static std::string Basename(const std::string &path) { + auto pos = path.find_last_of("/\\"); + if (pos == std::string::npos) { + return path; } - return idx; + return path.substr(pos + 1); } -template -static std::vector NpRepeat(const std::vector &v, - const std::vector ×) { - std::vector result; - for (size_t i = 0; i < times.size(); i++) { - for (int32_t n = 0; n < times[i]; n++) { - result.push_back(v[i]); - } +static int32_t ExtractIntegerAfterTag(const std::string &s, + const std::string &tag) { + auto pos = s.find(tag); + if (pos == std::string::npos) { + return -1; + } + + pos += tag.size(); + int32_t value = 0; + bool found = false; + while (pos < s.size() && std::isdigit(static_cast(s[pos]))) { + found = true; + value = value * 10 + (s[pos] - '0'); + ++pos; + } + + return found ? value : -1; +} + +static std::vector CropLastDim(const std::vector &src, + int32_t outer_dim, + int32_t src_last_dim, + int32_t dst_last_dim) { + std::vector ans(outer_dim * dst_last_dim); + for (int32_t i = 0; i != outer_dim; ++i) { + std::memcpy(ans.data() + i * dst_last_dim, + src.data() + i * src_last_dim, + dst_last_dim * sizeof(float)); } - return result; + + return ans; } -template -static std::vector Linspace(T a, T b, size_t N) { - if (N == 0) return {}; - if (N == 1) return {a}; - T h = (b - a) / static_cast(N - 1); - std::vector xs(N); - T val = a; - for (size_t i = 0; i < N; ++i) { - xs[i] = val; - val += h; +static std::vector PadLastDim(const std::vector &src, + int32_t outer_dim, int32_t src_last_dim, + int32_t dst_last_dim) { + std::vector ans(outer_dim * dst_last_dim, 0.0f); + for (int32_t i = 0; i != outer_dim; ++i) { + std::memcpy(ans.data() + i * dst_last_dim, + src.data() + i * src_last_dim, + src_last_dim * sizeof(float)); } - return xs; + + return ans; } -class OfflineTtsKokoroModelAxcl::Impl { - private: - int32_t max_seq_len_ = 0; - std::vector styles_; +static std::vector PadAlignment(const std::vector &src, + int32_t src_token_len, + int32_t src_frame_len, + int32_t dst_token_len, + int32_t dst_frame_len) { + std::vector ans(dst_token_len * dst_frame_len, 0.0f); + for (int32_t i = 0; i != src_token_len; ++i) { + std::memcpy(ans.data() + i * dst_frame_len, + src.data() + i * src_frame_len, + src_frame_len * sizeof(float)); + } + + return ans; +} + +static std::vector BuildPaddedTextMask(int32_t token_len, + int32_t token_bucket) { + std::vector mask(token_bucket, 0); + for (int32_t i = token_len; i < token_bucket; ++i) { + mask[i] = 1; + } + + return mask; +} + +struct ModelPaths { + std::string encoder; + std::string duration_predictor; + std::string text_encoder; + std::string f0n_shared; + std::string f0n_head; + std::string decoder_front; + std::string vocoder; +}; + +static ModelPaths ParseModelPaths(const std::string &model_str) { + std::vector files; + SplitStringToVector(model_str, ",", false, &files); + + ModelPaths paths; + for (const auto &f : files) { + if (!FileExists(f)) { + SHERPA_ONNX_LOGE("Model file '%s' does not exist", f.c_str()); + SHERPA_ONNX_EXIT(-1); + } + + std::string name = Basename(f); + if (name.find("duration_predictor") != std::string::npos) { + paths.duration_predictor = f; + } else if (name.find("text_encoder_token_") != std::string::npos) { + paths.text_encoder = f; + } else if (name.find("encoder_token_") != std::string::npos) { + paths.encoder = f; + } else if (name.find("f0n_shared_frame_") != std::string::npos) { + paths.f0n_shared = f; + } else if (name.find("f0n_head_frame_") != std::string::npos) { + paths.f0n_head = f; + } else if (name.find("decoder_front") != std::string::npos) { + paths.decoder_front = f; + } else if (name.find("vocoder") != std::string::npos) { + paths.vocoder = f; + } + } + + if (paths.encoder.empty() || paths.duration_predictor.empty() || + paths.text_encoder.empty() || paths.f0n_shared.empty() || + paths.f0n_head.empty() || paths.decoder_front.empty() || + paths.vocoder.empty()) { + SHERPA_ONNX_LOGE( + "For the new Axcl Kokoro pipeline, --kokoro-model must contain these " + "7 files in any order: encoder_token_*.axmodel, " + "duration_predictor.onnx, text_encoder_token_*_frame_*.axmodel, " + "f0n_shared_frame_*.axmodel, f0n_head_frame_*.axmodel, " + "decoder_front*.axmodel, vocoder*.axmodel. Given: %s", + model_str.c_str()); + SHERPA_ONNX_EXIT(-1); + } + + return paths; +} + +} // namespace +class OfflineTtsKokoroModelAxcl::Impl { public: explicit Impl(const OfflineTtsModelConfig &config) : config_(config), env_(ORT_LOGGING_LEVEL_ERROR), sess_opts_(GetSessionOptions(config)), allocator_{} { - Init(config); + Init(config, [](const std::string &filename) { return ReadFile(filename); }); } template @@ -105,32 +228,36 @@ class OfflineTtsKokoroModelAxcl::Impl { env_(ORT_LOGGING_LEVEL_ERROR), sess_opts_(GetSessionOptions(config)), allocator_{} { - Init(mgr, config); + Init(config, [mgr](const std::string &filename) { return ReadFile(mgr, filename); }); } const OfflineTtsKokoroModelMetaData &GetMetaData() const { return meta_data_; } - std::vector Run(std::vector input_ids, int64_t sid, + std::vector Run(const std::vector &input_ids, int64_t sid, float speed) { + std::lock_guard lock(mutex_); + if (input_ids.size() < 2) { return {}; } if (config_.kokoro.length_scale != 1 && speed == 1) { - speed = 1. / config_.kokoro.length_scale; + speed = 1.0f / config_.kokoro.length_scale; } int32_t phoneme_len = static_cast(input_ids.size()) - 2; - if (phoneme_len < 0) phoneme_len = 0; + if (phoneme_len < 0) { + phoneme_len = 0; + } - std::vector ref_s = LoadVoiceEmbedding(sid, phoneme_len); + std::vector ref_s = + LoadVoiceEmbedding(static_cast(sid), phoneme_len); std::vector audio; - bool success = RunModels(input_ids, ref_s, speed, audio); - if (!success) { - SHERPA_ONNX_LOGE("Run models failed!"); + if (!RunModels(input_ids, ref_s, speed, &audio)) { + SHERPA_ONNX_LOGE("Run Kokoro Axcl static pipeline failed"); return {}; } @@ -138,510 +265,480 @@ class OfflineTtsKokoroModelAxcl::Impl { } private: - void Init(const OfflineTtsModelConfig &config) { - std::string model_str = config.kokoro.model; - - std::vector model_files; - SplitStringToVector(model_str, ",", false, &model_files); - - std::vector model_paths; - - if (config_.provider == "axcl") { - // axcl provider requires exactly 4 model files - if (model_files.size() != 4) { - SHERPA_ONNX_LOGE( - "For axcl provider, model should contain exactly 4 files " - "separated by comma. Given %d files: %s", - static_cast(model_files.size()), model_str.c_str()); - SHERPA_ONNX_EXIT(-1); - } - - // Validate all 4 files exist - for (int32_t i = 0; i < 4; ++i) { - if (!FileExists(model_files[i])) { - SHERPA_ONNX_LOGE("Model file '%s' does not exist", - model_files[i].c_str()); - SHERPA_ONNX_EXIT(-1); - } - } - model_paths = model_files; - } else { + template + void Init(const OfflineTtsModelConfig &config, ReadBytes &&read_bytes) { + if (config_.provider != "axcl") { SHERPA_ONNX_LOGE( "This model only supports axcl provider. Please use provider=axcl"); SHERPA_ONNX_EXIT(-1); } - // Load ax models (indices 0, 1, 2) - for (int32_t i = 0; i < 3; ++i) { - axcl_models_.push_back(std::make_unique(model_paths[i])); - } + auto paths = ParseModelPaths(config.kokoro.model); - // Load onnx model (index 3) - auto buf4 = ReadFile(model_paths[3]); - model4_ = std::make_unique(env_, buf4.data(), buf4.size(), - sess_opts_); - GetInputNames(model4_.get(), &model4_input_names_, - &model4_input_names_ptr_); - GetOutputNames(model4_.get(), &model4_output_names_, - &model4_output_names_ptr_); + auto encoder_buf = read_bytes(paths.encoder); + encoder_model_ = std::make_unique(encoder_buf.data(), + encoder_buf.size()); - auto voices_buf = ReadFile(config.kokoro.voices); - LoadVoices(voices_buf.data(), voices_buf.size()); - } + auto text_encoder_buf = read_bytes(paths.text_encoder); + text_encoder_model_ = std::make_unique( + text_encoder_buf.data(), text_encoder_buf.size()); - template - void Init(Manager *mgr, const OfflineTtsModelConfig &config) { - std::string model_str = config.kokoro.model; - - std::vector model_files; - SplitStringToVector(model_str, ",", false, &model_files); - - std::vector model_paths; - - if (config_.provider == "axcl") { - // axcl provider requires exactly 4 model files - if (model_files.size() != 4) { - SHERPA_ONNX_LOGE( - "For axcl provider, model should contain exactly 4 files " - "separated by comma. Given %d files: %s", - static_cast(model_files.size()), model_str.c_str()); - SHERPA_ONNX_EXIT(-1); - } - model_paths = model_files; - } else { - SHERPA_ONNX_LOGE( - "This model only supports axcl provider. Please use provider=axcl"); - SHERPA_ONNX_EXIT(-1); - } + auto f0n_shared_buf = read_bytes(paths.f0n_shared); + f0n_shared_model_ = + std::make_unique(f0n_shared_buf.data(), f0n_shared_buf.size()); - // Load ax models (indices 0, 1, 2) - for (int32_t i = 0; i < 3; ++i) { - auto buf = ReadFile(mgr, model_paths[i]); - axcl_models_.push_back( - std::make_unique(buf.data(), buf.size())); - } + auto f0n_head_buf = read_bytes(paths.f0n_head); + f0n_head_model_ = + std::make_unique(f0n_head_buf.data(), f0n_head_buf.size()); + + auto decoder_front_buf = read_bytes(paths.decoder_front); + decoder_front_model_ = std::make_unique(decoder_front_buf.data(), + decoder_front_buf.size()); + + auto vocoder_buf = read_bytes(paths.vocoder); + vocoder_model_ = + std::make_unique(vocoder_buf.data(), vocoder_buf.size()); - // Load onnx model (index 3) - auto buf4 = ReadFile(mgr, model_paths[3]); - model4_ = std::make_unique(env_, buf4.data(), buf4.size(), - sess_opts_); - GetInputNames(model4_.get(), &model4_input_names_, - &model4_input_names_ptr_); - GetOutputNames(model4_.get(), &model4_output_names_, - &model4_output_names_ptr_); + auto duration_buf = read_bytes(paths.duration_predictor); + duration_predictor_ = std::make_unique( + env_, duration_buf.data(), duration_buf.size(), sess_opts_); + GetInputNames(duration_predictor_.get(), &duration_input_names_, + &duration_input_names_ptr_); + GetOutputNames(duration_predictor_.get(), &duration_output_names_, + &duration_output_names_ptr_); - auto voices_buf = ReadFile(mgr, config.kokoro.voices); + InitModelShapes(paths); + + auto voices_buf = read_bytes(config.kokoro.voices); LoadVoices(voices_buf.data(), voices_buf.size()); } - void LoadVoices(const char *voices_data, size_t voices_data_length) { - // Get metadata from model4 (ONNX model) - Ort::ModelMetadata meta_data = model4_->GetModelMetadata(); - Ort::AllocatorWithDefaultOptions allocator; - - SHERPA_ONNX_READ_META_DATA(meta_data_.sample_rate, "sample_rate"); - SHERPA_ONNX_READ_META_DATA_WITH_DEFAULT(meta_data_.version, "version", 1); - SHERPA_ONNX_READ_META_DATA(meta_data_.num_speakers, "n_speakers"); - SHERPA_ONNX_READ_META_DATA(meta_data_.has_espeak, "has_espeak"); - SHERPA_ONNX_READ_META_DATA_STR_WITH_DEFAULT(meta_data_.voice, "voice", - "en-us"); - SHERPA_ONNX_READ_META_DATA_VEC(style_dim_, "style_dim"); - - if (style_dim_.size() != 3) { - SHERPA_ONNX_LOGE("style_dim should be 3-d, given: %d", - static_cast(style_dim_.size())); + void InitModelShapes(const ModelPaths &paths) { + auto encoder_input_shape = encoder_model_->TensorShape("input_ids"); + auto encoder_mask_shape = encoder_model_->TensorShape("text_mask"); + auto text_encoder_input_shape = text_encoder_model_->TensorShape("input_ids"); + auto text_encoder_aln_shape = text_encoder_model_->TensorShape("pred_aln_trg"); + auto f0n_shared_input_shape = f0n_shared_model_->TensorShape("en"); + auto f0n_head_ref_shape = f0n_head_model_->TensorShape("ref_s"); + auto f0n_head_output_shape = f0n_head_model_->TensorShape("F0_pred"); + auto decoder_asr_shape = decoder_front_model_->TensorShape("asr"); + auto decoder_timbre_shape = decoder_front_model_->TensorShape("timbre"); + auto vocoder_har_shape = vocoder_model_->TensorShape("har"); + auto vocoder_output_shape = vocoder_model_->TensorShape("waveform"); + + if (encoder_input_shape.size() != 2 || encoder_mask_shape.size() != 2 || + text_encoder_input_shape.size() != 2 || text_encoder_aln_shape.size() != 3 || + f0n_shared_input_shape.size() != 3 || f0n_head_ref_shape.size() != 2 || + f0n_head_output_shape.size() != 2 || decoder_asr_shape.size() != 3 || + decoder_timbre_shape.size() != 2 || vocoder_har_shape.size() != 3) { + SHERPA_ONNX_LOGE("Unexpected Kokoro Axcl model shapes in: %s", + config_.kokoro.model.c_str()); SHERPA_ONNX_EXIT(-1); } - if (style_dim_[1] != 1) { - SHERPA_ONNX_LOGE("style_dim[1] should be 1, given: %d", style_dim_[1]); + token_bucket_ = encoder_input_shape[1]; + frame_bucket_ = text_encoder_aln_shape[2]; + pitch_bucket_ = f0n_head_output_shape[1]; + ref_s_dim_ = f0n_head_ref_shape[1]; + timbre_dim_ = decoder_timbre_shape[1]; + har_channels_ = vocoder_har_shape[1]; + har_frames_ = vocoder_har_shape[2]; + sample_bucket_ = GetNumElements(vocoder_output_shape); + + if (token_bucket_ != text_encoder_input_shape[1] || + token_bucket_ != text_encoder_aln_shape[1] || + frame_bucket_ != f0n_shared_input_shape[2] || + frame_bucket_ != decoder_asr_shape[2]) { + SHERPA_ONNX_LOGE("Kokoro Axcl bucket mismatch detected in: %s", + config_.kokoro.model.c_str()); SHERPA_ONNX_EXIT(-1); } - meta_data_.max_token_len = style_dim_[0]; + if (ref_s_dim_ <= 0 || timbre_dim_ <= 0 || timbre_dim_ > ref_s_dim_) { + SHERPA_ONNX_LOGE("Invalid Kokoro style dims. ref_s=%d, timbre=%d", + ref_s_dim_, timbre_dim_); + SHERPA_ONNX_EXIT(-1); + } - int32_t actual_num_floats = voices_data_length / sizeof(float); - int32_t expected_num_floats = - style_dim_[0] * style_dim_[2] * meta_data_.num_speakers; + if (har_channels_ != kStftNfft + 2) { + SHERPA_ONNX_LOGE("Expected har channels %d. Given: %d", kStftNfft + 2, + har_channels_); + SHERPA_ONNX_EXIT(-1); + } - if (actual_num_floats != expected_num_floats) { + if (sample_bucket_ % pitch_bucket_ != 0 || sample_bucket_ % frame_bucket_ != 0) { SHERPA_ONNX_LOGE( - "Corrupted voices file. Expected #floats: %d, actual: %d", - expected_num_floats, actual_num_floats); + "Unexpected Kokoro Axcl bucket relation. sample=%d pitch=%d frame=%d", + sample_bucket_, pitch_bucket_, frame_bucket_); SHERPA_ONNX_EXIT(-1); } - styles_ = std::vector( - reinterpret_cast(voices_data), - reinterpret_cast(voices_data) + actual_num_floats); + upsample_scale_ = sample_bucket_ / pitch_bucket_; + samples_per_frame_ = sample_bucket_ / frame_bucket_; + int32_t expected_har_frames = sample_bucket_ / kStftHop + 1; + if (har_frames_ != expected_har_frames) { + SHERPA_ONNX_LOGE("Expected har frames %d. Given: %d", expected_har_frames, + har_frames_); + SHERPA_ONNX_EXIT(-1); + } - auto shape1 = axcl_models_[0]->TensorShape("input_ids"); - max_seq_len_ = shape1[1]; + if (ExtractIntegerAfterTag(Basename(paths.encoder), "encoder_token_") != -1 && + ExtractIntegerAfterTag(Basename(paths.encoder), "encoder_token_") != + token_bucket_) { + SHERPA_ONNX_LOGE("Encoder filename bucket does not match model shape: %s", + paths.encoder.c_str()); + SHERPA_ONNX_EXIT(-1); + } } - std::vector LoadVoiceEmbedding(int32_t sid, int32_t phoneme_len) { - int32_t style_dim0 = style_dim_[0]; - int32_t style_dim1 = style_dim_[2]; - phoneme_len = std::max(phoneme_len, 0); - - std::vector ref_s(style_dim1); - if (phoneme_len < style_dim0) { - const float *p = styles_.data() + sid * style_dim0 * style_dim1 + - phoneme_len * style_dim1; - std::copy(p, p + style_dim1, ref_s.begin()); - } else { - int32_t idx = style_dim0 / 2; - const float *p = - styles_.data() + sid * style_dim0 * style_dim1 + idx * style_dim1; - std::copy(p, p + style_dim1, ref_s.begin()); + void LoadVoices(const char *voices_data, size_t voices_data_length) { + meta_data_.sample_rate = kKokoroSampleRate; + meta_data_.version = kKokoroVersion; + meta_data_.has_espeak = 1; + meta_data_.voice = config_.kokoro.lang.empty() ? kDefaultVoice : config_.kokoro.lang; + meta_data_.max_token_len = kKokoroVoiceStyleLength; + + style_dim_ = {kKokoroVoiceStyleLength, kKokoroStyleMiddleDim, ref_s_dim_}; + if (style_dim_.size() != kKokoroStyleRank || style_dim_[1] != 1) { + SHERPA_ONNX_LOGE("Unexpected Kokoro style_dim"); + SHERPA_ONNX_EXIT(-1); } - return ref_s; - } - bool RunModels(std::vector &input_ids, - const std::vector &ref_s, float speed, - std::vector &audio) { - int32_t actual_len = input_ids.size(); - if (actual_len > max_seq_len_) { + int32_t actual_num_floats = static_cast(voices_data_length / sizeof(float)); + int32_t floats_per_speaker = style_dim_[0] * style_dim_[2]; + if (floats_per_speaker <= 0 || actual_num_floats % floats_per_speaker != 0) { SHERPA_ONNX_LOGE( - "Input ids length %d exceeds max_seq_len %d, truncating.", actual_len, - max_seq_len_); - input_ids.resize(max_seq_len_); - actual_len = max_seq_len_; + "Corrupted voices file '%s'. Expected a multiple of %d floats. Got %d", + config_.kokoro.voices.c_str(), floats_per_speaker, actual_num_floats); + SHERPA_ONNX_EXIT(-1); } - int32_t padding_len = max_seq_len_ - actual_len; - if (padding_len > 0) { - std::vector padding(padding_len, 0); - input_ids.insert(input_ids.end(), padding.begin(), padding.end()); - } + meta_data_.num_speakers = actual_num_floats / floats_per_speaker; + styles_ = std::vector(reinterpret_cast(voices_data), + reinterpret_cast(voices_data) + + actual_num_floats); + } - int32_t actual_content_frames = 0; - int32_t total_frames = 0; - if (!InferenceSingleChunk(input_ids, ref_s, actual_len, speed, audio, - actual_content_frames, total_frames)) { - return false; + std::vector LoadVoiceEmbedding(int32_t sid, int32_t phoneme_len) const { + int32_t style_len = style_dim_[0]; + int32_t style_dim = style_dim_[2]; + + sid = std::max(sid, 0); + if (meta_data_.num_speakers > 0) { + sid = std::min(sid, meta_data_.num_speakers - 1); } - TrimAudioByContent(audio, actual_content_frames, total_frames, actual_len); + phoneme_len = std::max(phoneme_len, 0); - return true; + std::vector ref_s(style_dim); + int32_t index = phoneme_len < style_len ? phoneme_len : style_len / 2; + const float *src = styles_.data() + sid * style_len * style_dim + + index * style_dim; + std::copy(src, src + style_dim, ref_s.begin()); + return ref_s; } - bool InferenceSingleChunk(std::vector &input_ids, - const std::vector &ref_s, int32_t actual_len, - float speed, std::vector &audio, - int32_t &actual_content_frames, - int32_t &total_frames) { - bool is_doubled = false; - int32_t original_actual_len = actual_len; - - PrepareInputIds(input_ids, actual_len, is_doubled); - - std::vector input_ids_i32(input_ids.begin(), input_ids.end()); + std::vector BuildHar(const std::vector &f0_pred) const { + std::vector phase_coarse(pitch_bucket_ * kHarmonicDim, 0.0f); + std::array cumulative{}; + + for (int32_t t = 0; t != pitch_bucket_; ++t) { + float base_f0 = f0_pred[t]; + for (int32_t h = 0; h != kHarmonicDim; ++h) { + float harmonic = base_f0 * static_cast(h + 1); + float rad = std::fmod(harmonic / static_cast(kKokoroSampleRate), 1.0f); + if (rad < 0) { + rad += 1.0f; + } + cumulative[h] += rad; + phase_coarse[t * kHarmonicDim + h] = + cumulative[h] * kTwoPi * static_cast(upsample_scale_); + } + } - std::vector input_lengths; - std::vector text_mask; - ComputeExternalPreprocessing(input_ids, actual_len, input_lengths, - text_mask); + std::vector sine_merge(sample_bucket_); + std::mt19937 rng(0); + std::normal_distribution normal(0.0f, 1.0f); + + for (int32_t sample_index = 0; sample_index != sample_bucket_; ++sample_index) { + int32_t coarse_index = sample_index / upsample_scale_; + float coarse_f0 = f0_pred[coarse_index]; + float uv = coarse_f0 > kVoicedThreshold ? 1.0f : 0.0f; + float noise_amp = uv * kNoiseStd + (1.0f - uv) * kUnvoicedNoiseScale; + + float src = (static_cast(sample_index) + 0.5f) / + static_cast(upsample_scale_) - + 0.5f; + if (src < 0.0f) { + src = 0.0f; + } - std::vector text_mask_u8(text_mask.begin(), text_mask.end()); + int32_t left = static_cast(std::floor(src)); + int32_t right = std::min(left + 1, pitch_bucket_ - 1); + float w = src - static_cast(left); + if (left >= pitch_bucket_ - 1) { + left = pitch_bucket_ - 1; + right = left; + w = 0.0f; + } - axcl_models_[0]->SetInputTensorData("input_ids", input_ids_i32.data(), - input_ids_i32.size()); - axcl_models_[0]->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); - axcl_models_[0]->SetInputTensorData("text_mask", text_mask_u8.data(), - text_mask_u8.size()); + float merged = kHarmonicMergeBias; + for (int32_t h = 0; h != kHarmonicDim; ++h) { + float phase_left = phase_coarse[left * kHarmonicDim + h]; + float phase_right = phase_coarse[right * kHarmonicDim + h]; + float phase = phase_left * (1.0f - w) + phase_right * w; + float sine = std::sin(phase) * kSineAmp; + float noise = noise_amp * normal(rng); + float sine_wave = sine * uv + noise; + merged += sine_wave * kHarmonicMergeWeights[h]; + } - if (!axcl_models_[0]->Run()) { - return false; + sine_merge[sample_index] = std::tanh(merged); } - std::vector duration = - axcl_models_[0]->GetOutputTensorData("duration"); - std::vector d = axcl_models_[0]->GetOutputTensorData("d"); + std::vector padded(sample_bucket_ + kStftNfft, 0.0f); + int32_t pad = kStftNfft / 2; + std::copy(sine_merge.begin(), sine_merge.end(), padded.begin() + pad); + std::fill(padded.begin(), padded.begin() + pad, sine_merge.front()); + std::fill(padded.end() - pad, padded.end(), sine_merge.back()); - std::vector pred_dur; - ProcessDuration(duration, actual_len, speed, pred_dur, total_frames); + std::array window{}; + std::array, kStftNfft / 2 + 1> cos_table{}; + std::array, kStftNfft / 2 + 1> sin_table{}; - std::vector pred_aln_trg = - CreateAlignmentMatrix(pred_dur, total_frames); + for (int32_t n = 0; n != kStftNfft; ++n) { + window[n] = 0.5f - 0.5f * std::cos(kTwoPi * n / static_cast(kStftNfft)); + } - int32_t d_shape_1 = max_seq_len_; - int32_t d_shape_2 = - 640; // style_dim_[2] essentially, or based on d_shape_[2] - std::vector en(total_frames * d_shape_2, 0.0f); + for (int32_t k = 0; k != kStftNfft / 2 + 1; ++k) { + for (int32_t n = 0; n != kStftNfft; ++n) { + float angle = kTwoPi * k * n / static_cast(kStftNfft); + cos_table[k][n] = std::cos(angle); + sin_table[k][n] = -std::sin(angle); + } + } - for (int32_t j = 0; j < d_shape_2; ++j) { - for (int32_t k = 0; k < total_frames; ++k) { - float sum = 0.0f; - for (int32_t i = 0; i < max_seq_len_; ++i) { - sum += d[i * d_shape_2 + j] * pred_aln_trg[i * total_frames + k]; + std::vector har(har_channels_ * har_frames_); + for (int32_t frame = 0; frame != har_frames_; ++frame) { + const float *segment = padded.data() + frame * kStftHop; + for (int32_t k = 0; k != kStftNfft / 2 + 1; ++k) { + float real = 0.0f; + float imag = 0.0f; + for (int32_t n = 0; n != kStftNfft; ++n) { + float value = segment[n] * window[n]; + real += value * cos_table[k][n]; + imag += value * sin_table[k][n]; } - en[j * total_frames + k] = sum; + + har[k * har_frames_ + frame] = + std::sqrt(real * real + imag * imag + 1e-14f); + float phase = std::atan2(imag, real); + if (imag == 0.0f && real < 0.0f) { + phase = kPi; + } + har[(k + kStftNfft / 2 + 1) * har_frames_ + frame] = phase; } } - std::vector text_mask_float(text_mask.begin(), text_mask.end()); + return har; + } + + bool RunModels(const std::vector &input_ids, + const std::vector &ref_s, float speed, + std::vector *audio) { + int32_t token_length = static_cast(input_ids.size()); + if (token_length > token_bucket_) { + SHERPA_ONNX_LOGE( + "Token length %d exceeds Axcl token bucket %d. Use a shorter text " + "or export a larger static bucket.", + token_length, token_bucket_); + return false; + } + + std::vector padded_input_ids(token_bucket_, 0); + for (int32_t i = 0; i != token_length; ++i) { + padded_input_ids[i] = static_cast(input_ids[i]); + } - axcl_models_[1]->SetInputTensorData("en", en.data(), en.size()); - axcl_models_[1]->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); - axcl_models_[1]->SetInputTensorData("input_ids", input_ids_i32.data(), - input_ids_i32.size()); - axcl_models_[1]->SetInputTensorData("text_mask", text_mask_float.data(), - text_mask_float.size()); - axcl_models_[1]->SetInputTensorData("pred_aln_trg", pred_aln_trg.data(), - pred_aln_trg.size()); + std::vector padded_text_mask = + BuildPaddedTextMask(token_length, token_bucket_); - if (!axcl_models_[1]->Run()) { + if (!encoder_model_->SetInputTensorData("input_ids", padded_input_ids.data(), + padded_input_ids.size()) || + !encoder_model_->SetInputTensorData("text_mask", padded_text_mask.data(), + padded_text_mask.size()) || + !encoder_model_->Run()) { return false; } - std::vector f0_pred = - axcl_models_[1]->GetOutputTensorData("F0_pred"); - std::vector n_pred = axcl_models_[1]->GetOutputTensorData("N_pred"); - std::vector asr = axcl_models_[1]->GetOutputTensorData("asr"); + std::vector padded_d_en = encoder_model_->GetOutputTensorData("d_en"); + std::vector d_en = + CropLastDim(padded_d_en, 512, token_bucket_, token_length); auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); - std::array f0_shape = {1, static_cast(f0_pred.size())}; - Ort::Value f0_tensor = - Ort::Value::CreateTensor(memory_info, f0_pred.data(), f0_pred.size(), - f0_shape.data(), f0_shape.size()); - - auto out4 = model4_->Run({}, model4_input_names_ptr_.data(), &f0_tensor, 1, - model4_output_names_ptr_.data(), - model4_output_names_ptr_.size()); - float *p_har = out4[0].GetTensorMutableData(); - auto har_shape = out4[0].GetTensorTypeAndShapeInfo().GetShape(); - int64_t har_num_elements = 1; - for (auto d : har_shape) { - har_num_elements *= d; - } - std::vector har(p_har, p_har + har_num_elements); - - axcl_models_[2]->SetInputTensorData("asr", asr.data(), asr.size()); - axcl_models_[2]->SetInputTensorData("F0_pred", f0_pred.data(), - f0_pred.size()); - axcl_models_[2]->SetInputTensorData("N_pred", n_pred.data(), n_pred.size()); - axcl_models_[2]->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()); - axcl_models_[2]->SetInputTensorData("har", har.data(), har.size()); - - if (!axcl_models_[2]->Run()) { + + std::array d_en_shape = {1, 512, token_length}; + std::array ref_s_shape = {1, ref_s_dim_}; + std::array input_lengths_shape = {1}; + std::array duration_text_mask_shape = {1, token_length}; + std::array speed_shape = {1}; + + std::array input_lengths = {token_length}; + std::unique_ptr duration_text_mask(new bool[token_length]()); + + Ort::Value d_en_tensor = Ort::Value::CreateTensor( + memory_info, d_en.data(), d_en.size(), d_en_shape.data(), d_en_shape.size()); + Ort::Value ref_s_tensor = Ort::Value::CreateTensor( + memory_info, const_cast(ref_s.data()), ref_s.size(), + ref_s_shape.data(), ref_s_shape.size()); + Ort::Value input_lengths_tensor = Ort::Value::CreateTensor( + memory_info, input_lengths.data(), input_lengths.size(), + input_lengths_shape.data(), input_lengths_shape.size()); + Ort::Value duration_text_mask_tensor = Ort::Value::CreateTensor( + memory_info, duration_text_mask.get(), token_length, + duration_text_mask_shape.data(), duration_text_mask_shape.size()); + Ort::Value speed_tensor = Ort::Value::CreateTensor( + memory_info, &speed, 1, speed_shape.data(), speed_shape.size()); + + std::array duration_inputs = { + std::move(d_en_tensor), std::move(ref_s_tensor), + std::move(input_lengths_tensor), std::move(duration_text_mask_tensor), + std::move(speed_tensor)}; + + auto duration_outputs = duration_predictor_->Run( + {}, duration_input_names_ptr_.data(), duration_inputs.data(), + duration_inputs.size(), duration_output_names_ptr_.data(), + duration_output_names_ptr_.size()); + + auto pred_aln_shape = duration_outputs[2].GetTensorTypeAndShapeInfo().GetShape(); + auto en_shape = duration_outputs[3].GetTensorTypeAndShapeInfo().GetShape(); + if (pred_aln_shape.size() != 3 || en_shape.size() != 3) { + SHERPA_ONNX_LOGE("Unexpected duration predictor output shapes"); return false; } - std::vector x_out = axcl_models_[2]->GetOutputTensorData("x"); - PostprocessXToAudio(x_out, 23041, audio); - - if (is_doubled) { - actual_content_frames = std::accumulate( - pred_dur.begin(), pred_dur.begin() + original_actual_len, 0); - audio.erase(audio.begin() + audio.size() / 2, audio.end()); - total_frames = total_frames / 2; - } else { - actual_content_frames = - std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); + int32_t frame_length = static_cast(pred_aln_shape[2]); + if (frame_length > frame_bucket_) { + SHERPA_ONNX_LOGE( + "Frame length %d exceeds Axcl frame bucket %d. Use a shorter text, " + "increase speed, or export a larger static bucket.", + frame_length, frame_bucket_); + return false; } - return true; - } - void ProcessDuration(const std::vector &duration, int32_t actual_len, - float speed, std::vector &pred_dur, - int32_t &total_frames) { - std::vector pred_dur_original(actual_len, 0); - std::vector duration_processed = Sigmoid(duration); - - int32_t dur_shape_2 = 50; - for (int32_t i = 0; i < actual_len; i++) { - float sum = 0; - for (int32_t n = 0; n < dur_shape_2; n++) { - sum += duration_processed[i * dur_shape_2 + n]; - } - sum /= speed; - pred_dur_original[i] = - static_cast(std::max(1.f, std::round(sum))); - } - - std::vector pred_dur_padding(max_seq_len_ - actual_len, 0); - pred_dur = pred_dur_original; - pred_dur.insert(pred_dur.end(), pred_dur_padding.begin(), - pred_dur_padding.end()); - - int32_t fixed_total_frames = max_seq_len_ * 2; - int32_t actual_frames = - std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); - int32_t diff = fixed_total_frames - actual_frames; - - if (diff < 0) { - auto indices = Argsort(pred_dur, actual_len, true); - int32_t decreased = 0; - for (auto idx : indices) { - if (pred_dur[idx] > 1 && decreased < std::abs(diff)) { - pred_dur[idx]--; - decreased++; - } - if (decreased >= std::abs(diff)) break; - } + const float *pred_aln_ptr = duration_outputs[2].GetTensorData(); + int32_t pred_aln_elems = static_cast( + duration_outputs[2].GetTensorTypeAndShapeInfo().GetElementCount()); + std::vector pred_aln_trg(pred_aln_ptr, pred_aln_ptr + pred_aln_elems); + std::vector padded_alignment = PadAlignment( + pred_aln_trg, token_length, frame_length, token_bucket_, frame_bucket_); + + if (!text_encoder_model_->SetInputTensorData("input_ids", padded_input_ids.data(), + padded_input_ids.size()) || + !text_encoder_model_->SetInputTensorData("pred_aln_trg", + padded_alignment.data(), + padded_alignment.size()) || + !text_encoder_model_->SetInputTensorData("text_mask", + padded_text_mask.data(), + padded_text_mask.size()) || + !text_encoder_model_->Run()) { + return false; } - actual_frames = - std::accumulate(pred_dur.begin(), pred_dur.begin() + actual_len, 0); - int32_t remaining_frames = fixed_total_frames - actual_frames; - int32_t padding_len = max_seq_len_ - actual_len; - - if (remaining_frames > 0 && padding_len > 0) { - int32_t frames_per_padding = remaining_frames / padding_len; - int32_t remainder = remaining_frames % padding_len; + std::vector asr = text_encoder_model_->GetOutputTensorData("asr"); - for (int32_t i = actual_len; i < static_cast(pred_dur.size()); - i++) - pred_dur[i] = frames_per_padding; + const float *en_ptr = duration_outputs[3].GetTensorData(); + int32_t en_elems = static_cast( + duration_outputs[3].GetTensorTypeAndShapeInfo().GetElementCount()); + std::vector en(en_ptr, en_ptr + en_elems); + std::vector padded_en = PadLastDim(en, 640, frame_length, frame_bucket_); - if (remainder > 0) { - for (int32_t i = actual_len; i < actual_len + remainder; i++) - pred_dur[i] += 1; - } + if (!f0n_shared_model_->SetInputTensorData("en", padded_en.data(), + padded_en.size()) || + !f0n_shared_model_->Run()) { + return false; } - total_frames = std::accumulate(pred_dur.begin(), pred_dur.end(), 0); - } - - std::vector CreateAlignmentMatrix(const std::vector &pred_dur, - int32_t total_frames) { - std::vector seq_range(max_seq_len_); - std::iota(seq_range.begin(), seq_range.end(), 0); - auto indices = NpRepeat(seq_range, pred_dur); - - std::vector pred_aln_trg(max_seq_len_ * total_frames, 0.0f); - if (!indices.empty()) { - int32_t col = 0; - for (auto i : indices) { - pred_aln_trg[i * total_frames + col] = 1.0f; - col++; - } + std::vector shared = f0n_shared_model_->GetOutputTensorData("shared"); + if (!f0n_head_model_->SetInputTensorData("shared", shared.data(), shared.size()) || + !f0n_head_model_->SetInputTensorData("ref_s", ref_s.data(), ref_s.size()) || + !f0n_head_model_->Run()) { + return false; } - return pred_aln_trg; - } - - void PostprocessXToAudio(const std::vector &x, int32_t num_frames, - std::vector &audio) { - int32_t n_fft = 20; - int32_t hop_length = 5; - int32_t half_n_fft = n_fft / 2 + 1; - - std::vector spec_part(half_n_fft * num_frames); - std::vector phase_part(half_n_fft * num_frames); - std::vector cos_part(half_n_fft * num_frames); - spec_part.assign(x.begin(), x.begin() + half_n_fft * num_frames); - phase_part.assign(x.begin() + half_n_fft * num_frames, x.end()); - - for (int32_t i = 0; i < half_n_fft * num_frames; i++) { - spec_part[i] = std::exp(spec_part[i]); - phase_part[i] = std::sin(phase_part[i]); - cos_part[i] = std::sqrt( - 1.f - std::max(0.f, std::min(phase_part[i] * phase_part[i], 1.0f))); - } - - knf::StftResult stft_result; - stft_result.num_frames = num_frames; - stft_result.real.resize(half_n_fft * num_frames); - stft_result.imag.resize(half_n_fft * num_frames); - - for (int32_t i = 0; i < half_n_fft; i++) { - for (int32_t n = 0; n < num_frames; n++) { - float spec = spec_part[i * num_frames + n]; - float real_part = spec * cos_part[i * num_frames + n]; - float imag_part = spec * phase_part[i * num_frames + n]; - - // StftResult layout is [frame_index * half_n_fft + bin] - stft_result.real[n * half_n_fft + i] = real_part; - stft_result.imag[n * half_n_fft + i] = imag_part; - } + std::vector f0_pred = f0n_head_model_->GetOutputTensorData("F0_pred"); + std::vector n_pred = f0n_head_model_->GetOutputTensorData("N_pred"); + std::vector timbre(ref_s.begin(), ref_s.begin() + timbre_dim_); + + if (!decoder_front_model_->SetInputTensorData("asr", asr.data(), asr.size()) || + !decoder_front_model_->SetInputTensorData("F0_pred", f0_pred.data(), + f0_pred.size()) || + !decoder_front_model_->SetInputTensorData("N_pred", n_pred.data(), + n_pred.size()) || + !decoder_front_model_->SetInputTensorData("timbre", timbre.data(), + timbre.size()) || + !decoder_front_model_->Run()) { + return false; } - knf::StftConfig stft_config; - stft_config.n_fft = n_fft; - stft_config.hop_length = hop_length; - stft_config.win_length = n_fft; - stft_config.center = true; - stft_config.normalized = false; - stft_config.window_type = "hann"; - stft_config.pad_mode = "reflect"; + std::vector decoder_state = + decoder_front_model_->GetOutputTensorData("decoder_state"); + std::vector har = BuildHar(f0_pred); - knf::IStft istft(stft_config); - audio = istft.Compute(stft_result); - } - - void TrimAudioByContent(std::vector &audio, - int32_t actual_content_frames, int32_t total_frames, - int32_t actual_len) { - int32_t padding_len = max_seq_len_ - actual_len; - if (padding_len > 0) { - float content_ratio = actual_content_frames * 1.0f / total_frames; - int32_t audio_len_to_keep = - static_cast(audio.size() * content_ratio); - audio.resize(audio_len_to_keep); + if (!vocoder_model_->SetInputTensorData("decoder_state", decoder_state.data(), + decoder_state.size()) || + !vocoder_model_->SetInputTensorData("timbre", timbre.data(), + timbre.size()) || + !vocoder_model_->SetInputTensorData("har", har.data(), har.size()) || + !vocoder_model_->Run()) { + return false; } - } - - void PrepareInputIds(std::vector &input_ids, int32_t &actual_len, - bool &is_doubled) { - is_doubled = false; - int32_t original_actual_len = actual_len; - - if (actual_len <= 32) { // DOUBLE_INPUT_THRESHOLD - is_doubled = true; - std::vector valid_content(input_ids.begin(), - input_ids.begin() + actual_len); - std::vector input_ids_doubled; - input_ids_doubled.insert(input_ids_doubled.end(), valid_content.begin(), - valid_content.end()); - input_ids_doubled.insert(input_ids_doubled.end(), valid_content.begin(), - valid_content.end()); - - int32_t padding_len = max_seq_len_ - 2 * actual_len; - if (padding_len > 0) { - std::vector padding(padding_len, 0); - input_ids_doubled.insert(input_ids_doubled.end(), padding.begin(), - padding.end()); - } else { - input_ids_doubled.resize(max_seq_len_); - } - input_ids = input_ids_doubled; - actual_len = std::min(original_actual_len * 2, max_seq_len_); + *audio = vocoder_model_->GetOutputTensorData("waveform"); + int32_t sample_length = frame_length * samples_per_frame_; + if (sample_length < static_cast(audio->size())) { + audio->resize(sample_length); } - } - void ComputeExternalPreprocessing(const std::vector &input_ids, - int32_t actual_len, - std::vector &input_lengths, - std::vector &text_mask) { - input_lengths = std::vector{actual_len}; - text_mask.resize(max_seq_len_); - for (int32_t i = 0; i < max_seq_len_; i++) { - text_mask[i] = (i >= actual_len) ? 1 : 0; - } + return true; } private: - OfflineTtsModelConfig config_; - OfflineTtsKokoroModelMetaData meta_data_; - std::vector style_dim_; + mutable std::mutex mutex_; + OfflineTtsModelConfig config_; Ort::Env env_; Ort::SessionOptions sess_opts_; Ort::AllocatorWithDefaultOptions allocator_; - std::vector> axcl_models_; + std::unique_ptr encoder_model_; + std::unique_ptr text_encoder_model_; + std::unique_ptr f0n_shared_model_; + std::unique_ptr f0n_head_model_; + std::unique_ptr decoder_front_model_; + std::unique_ptr vocoder_model_; + + std::unique_ptr duration_predictor_; + std::vector duration_input_names_; + std::vector duration_input_names_ptr_; + std::vector duration_output_names_; + std::vector duration_output_names_ptr_; + + OfflineTtsKokoroModelMetaData meta_data_; + std::vector style_dim_; + std::vector styles_; - std::unique_ptr model4_; - std::vector model4_input_names_; - std::vector model4_input_names_ptr_; - std::vector model4_output_names_; - std::vector model4_output_names_ptr_; + int32_t token_bucket_ = 0; + int32_t frame_bucket_ = 0; + int32_t pitch_bucket_ = 0; + int32_t sample_bucket_ = 0; + int32_t ref_s_dim_ = 0; + int32_t timbre_dim_ = 0; + int32_t har_channels_ = 0; + int32_t har_frames_ = 0; + int32_t upsample_scale_ = 0; + int32_t samples_per_frame_ = 0; }; OfflineTtsKokoroModelAxcl::~OfflineTtsKokoroModelAxcl() = default; @@ -655,9 +752,8 @@ OfflineTtsKokoroModelAxcl::OfflineTtsKokoroModelAxcl( Manager *mgr, const OfflineTtsModelConfig &config) : impl_(std::make_unique(mgr, config)) {} -std::vector OfflineTtsKokoroModelAxcl::Run(const std::vector &x, - int64_t sid, - float speed) const { +std::vector OfflineTtsKokoroModelAxcl::Run( + const std::vector &x, int64_t sid, float speed) const { return impl_->Run(x, sid, speed); } From 3dd5b298502f5de1b02554bf4711190cb842cfb0 Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Tue, 7 Apr 2026 10:52:06 +0800 Subject: [PATCH 09/10] Add BSP SDK and model files to .gitignore --- .gitignore | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.gitignore b/.gitignore index 25eee9ab08..6276abd36b 100755 --- a/.gitignore +++ b/.gitignore @@ -177,3 +177,9 @@ sherpa-onnx-fire-red-asr2-ctc-zh_en-int8-2026-02-25 non-streaming-fire-red-asr-ctc-decode-files sherpa-onnx-moonshine-*-quantized-2026-02-27 sherpa-onnx-supertonic-tts-int8-2026-03-06 + +ax620e_bsp_sdk-2.0.0_P7 +ax650n_bsp_sdk-1.45.0_p39 +aarch64-linux-gnu-g* +kokoro-multi-lang-v1_1 +kokoro-v1_1-zh.pth From 6823405036de2d4140b81f0615a4142d23b2e90b Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Tue, 7 Apr 2026 14:29:26 +0800 Subject: [PATCH 10/10] Enhance Axcl and Axera models with new tensor shape handling and update vocoder paths --- .../axcl/offline-tts-kokoro-model-axcl.cc | 147 +++++++++++++++--- .../axera/offline-tts-kokoro-model-axera.cc | 147 +++++++++++++++--- 2 files changed, 248 insertions(+), 46 deletions(-) diff --git a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc index 8afd50202b..f85d0527f2 100644 --- a/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc +++ b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc @@ -74,6 +74,19 @@ static int32_t GetNumElements(const std::vector &shape) { return ans; } +static int64_t GetNumElements(const std::vector &shape) { + if (shape.empty()) { + return 0; + } + + int64_t ans = 1; + for (auto dim : shape) { + ans *= dim; + } + + return ans; +} + static std::string Basename(const std::string &path) { auto pos = path.find_last_of("/\\"); if (pos == std::string::npos) { @@ -154,6 +167,28 @@ static std::vector BuildPaddedTextMask(int32_t token_len, return mask; } +static std::vector GetOrtTensorShape(Ort::Session *sess, + const std::string &name, + bool is_input) { + Ort::AllocatorWithDefaultOptions allocator; + size_t count = is_input ? sess->GetInputCount() : sess->GetOutputCount(); + for (size_t i = 0; i != count; ++i) { + Ort::AllocatedStringPtr tensor_name = + is_input ? sess->GetInputNameAllocated(i, allocator) + : sess->GetOutputNameAllocated(i, allocator); + + if (name == tensor_name.get()) { + auto type_info = + is_input ? sess->GetInputTypeInfo(i) : sess->GetOutputTypeInfo(i); + return type_info.GetTensorTypeAndShapeInfo().GetShape(); + } + } + + SHERPA_ONNX_LOGE("Found no %s tensor with name: '%s'", + is_input ? "input" : "output", name.c_str()); + return {}; +} + struct ModelPaths { std::string encoder; std::string duration_predictor; @@ -161,7 +196,8 @@ struct ModelPaths { std::string f0n_shared; std::string f0n_head; std::string decoder_front; - std::string vocoder; + std::string vocoder_core; + std::string vocoder_tail; }; static ModelPaths ParseModelPaths(const std::string &model_str) { @@ -188,21 +224,24 @@ static ModelPaths ParseModelPaths(const std::string &model_str) { paths.f0n_head = f; } else if (name.find("decoder_front") != std::string::npos) { paths.decoder_front = f; - } else if (name.find("vocoder") != std::string::npos) { - paths.vocoder = f; + } else if (name.find("vocoder_tail") != std::string::npos) { + paths.vocoder_tail = f; + } else if (name.find("vocoder_core") != std::string::npos) { + paths.vocoder_core = f; } } if (paths.encoder.empty() || paths.duration_predictor.empty() || paths.text_encoder.empty() || paths.f0n_shared.empty() || paths.f0n_head.empty() || paths.decoder_front.empty() || - paths.vocoder.empty()) { + paths.vocoder_core.empty() || paths.vocoder_tail.empty()) { SHERPA_ONNX_LOGE( "For the new Axcl Kokoro pipeline, --kokoro-model must contain these " - "7 files in any order: encoder_token_*.axmodel, " + "8 files in any order: encoder_token_*.axmodel, " "duration_predictor.onnx, text_encoder_token_*_frame_*.axmodel, " "f0n_shared_frame_*.axmodel, f0n_head_frame_*.axmodel, " - "decoder_front*.axmodel, vocoder*.axmodel. Given: %s", + "decoder_front*.axmodel, vocoder_core*.axmodel, " + "vocoder_tail*.onnx. Given: %s", model_str.c_str()); SHERPA_ONNX_EXIT(-1); } @@ -295,9 +334,17 @@ class OfflineTtsKokoroModelAxcl::Impl { decoder_front_model_ = std::make_unique(decoder_front_buf.data(), decoder_front_buf.size()); - auto vocoder_buf = read_bytes(paths.vocoder); - vocoder_model_ = - std::make_unique(vocoder_buf.data(), vocoder_buf.size()); + auto vocoder_core_buf = read_bytes(paths.vocoder_core); + vocoder_core_model_ = std::make_unique(vocoder_core_buf.data(), + vocoder_core_buf.size()); + + auto vocoder_tail_buf = read_bytes(paths.vocoder_tail); + vocoder_tail_ = std::make_unique( + env_, vocoder_tail_buf.data(), vocoder_tail_buf.size(), sess_opts_); + GetInputNames(vocoder_tail_.get(), &vocoder_tail_input_names_, + &vocoder_tail_input_names_ptr_); + GetOutputNames(vocoder_tail_.get(), &vocoder_tail_output_names_, + &vocoder_tail_output_names_ptr_); auto duration_buf = read_bytes(paths.duration_predictor); duration_predictor_ = std::make_unique( @@ -323,14 +370,27 @@ class OfflineTtsKokoroModelAxcl::Impl { auto f0n_head_output_shape = f0n_head_model_->TensorShape("F0_pred"); auto decoder_asr_shape = decoder_front_model_->TensorShape("asr"); auto decoder_timbre_shape = decoder_front_model_->TensorShape("timbre"); - auto vocoder_har_shape = vocoder_model_->TensorShape("har"); - auto vocoder_output_shape = vocoder_model_->TensorShape("waveform"); + auto vocoder_core_state_shape = + vocoder_core_model_->TensorShape("decoder_state"); + auto vocoder_core_timbre_shape = vocoder_core_model_->TensorShape("timbre"); + auto vocoder_core_har_shape = vocoder_core_model_->TensorShape("har"); + auto vocoder_core_hidden_shape = + vocoder_core_model_->TensorShape("vocoder_hidden"); + auto vocoder_tail_hidden_shape = + GetOrtTensorShape(vocoder_tail_.get(), "vocoder_hidden", true); + auto vocoder_tail_output_shape = + GetOrtTensorShape(vocoder_tail_.get(), "waveform", false); if (encoder_input_shape.size() != 2 || encoder_mask_shape.size() != 2 || text_encoder_input_shape.size() != 2 || text_encoder_aln_shape.size() != 3 || f0n_shared_input_shape.size() != 3 || f0n_head_ref_shape.size() != 2 || f0n_head_output_shape.size() != 2 || decoder_asr_shape.size() != 3 || - decoder_timbre_shape.size() != 2 || vocoder_har_shape.size() != 3) { + decoder_timbre_shape.size() != 2 || vocoder_core_state_shape.size() != 3 || + vocoder_core_timbre_shape.size() != 2 || + vocoder_core_har_shape.size() != 3 || + vocoder_core_hidden_shape.size() != 3 || + vocoder_tail_hidden_shape.size() != 3 || + vocoder_tail_output_shape.size() != 1) { SHERPA_ONNX_LOGE("Unexpected Kokoro Axcl model shapes in: %s", config_.kokoro.model.c_str()); SHERPA_ONNX_EXIT(-1); @@ -341,9 +401,11 @@ class OfflineTtsKokoroModelAxcl::Impl { pitch_bucket_ = f0n_head_output_shape[1]; ref_s_dim_ = f0n_head_ref_shape[1]; timbre_dim_ = decoder_timbre_shape[1]; - har_channels_ = vocoder_har_shape[1]; - har_frames_ = vocoder_har_shape[2]; - sample_bucket_ = GetNumElements(vocoder_output_shape); + har_channels_ = vocoder_core_har_shape[1]; + har_frames_ = vocoder_core_har_shape[2]; + vocoder_hidden_channels_ = vocoder_core_hidden_shape[1]; + vocoder_hidden_frames_ = vocoder_core_hidden_shape[2]; + sample_bucket_ = static_cast(GetNumElements(vocoder_tail_output_shape)); if (token_bucket_ != text_encoder_input_shape[1] || token_bucket_ != text_encoder_aln_shape[1] || @@ -360,6 +422,15 @@ class OfflineTtsKokoroModelAxcl::Impl { SHERPA_ONNX_EXIT(-1); } + if (vocoder_core_state_shape != decoder_front_model_->TensorShape("decoder_state") || + vocoder_core_timbre_shape[1] != timbre_dim_ || + vocoder_tail_hidden_shape[1] != vocoder_hidden_channels_ || + vocoder_tail_hidden_shape[2] != vocoder_hidden_frames_) { + SHERPA_ONNX_LOGE("Kokoro Axcl vocoder_core/tail shape mismatch in: %s", + config_.kokoro.model.c_str()); + SHERPA_ONNX_EXIT(-1); + } + if (har_channels_ != kStftNfft + 2) { SHERPA_ONNX_LOGE("Expected har channels %d. Given: %d", kStftNfft + 2, har_channels_); @@ -686,16 +757,39 @@ class OfflineTtsKokoroModelAxcl::Impl { decoder_front_model_->GetOutputTensorData("decoder_state"); std::vector har = BuildHar(f0_pred); - if (!vocoder_model_->SetInputTensorData("decoder_state", decoder_state.data(), - decoder_state.size()) || - !vocoder_model_->SetInputTensorData("timbre", timbre.data(), - timbre.size()) || - !vocoder_model_->SetInputTensorData("har", har.data(), har.size()) || - !vocoder_model_->Run()) { + if (!vocoder_core_model_->SetInputTensorData("decoder_state", + decoder_state.data(), + decoder_state.size()) || + !vocoder_core_model_->SetInputTensorData("timbre", timbre.data(), + timbre.size()) || + !vocoder_core_model_->SetInputTensorData("har", har.data(), har.size()) || + !vocoder_core_model_->Run()) { + return false; + } + + std::vector vocoder_hidden = + vocoder_core_model_->GetOutputTensorData("vocoder_hidden"); + std::array vocoder_hidden_shape = { + 1, vocoder_hidden_channels_, vocoder_hidden_frames_}; + Ort::Value vocoder_hidden_tensor = Ort::Value::CreateTensor( + memory_info, vocoder_hidden.data(), vocoder_hidden.size(), + vocoder_hidden_shape.data(), vocoder_hidden_shape.size()); + std::array vocoder_tail_inputs = { + std::move(vocoder_hidden_tensor)}; + + auto vocoder_tail_outputs = vocoder_tail_->Run( + {}, vocoder_tail_input_names_ptr_.data(), vocoder_tail_inputs.data(), + vocoder_tail_inputs.size(), vocoder_tail_output_names_ptr_.data(), + vocoder_tail_output_names_ptr_.size()); + if (vocoder_tail_outputs.empty()) { + SHERPA_ONNX_LOGE("Unexpected empty vocoder_tail outputs"); return false; } - *audio = vocoder_model_->GetOutputTensorData("waveform"); + const float *waveform_ptr = vocoder_tail_outputs[0].GetTensorData(); + int32_t waveform_elems = static_cast( + vocoder_tail_outputs[0].GetTensorTypeAndShapeInfo().GetElementCount()); + *audio = std::vector(waveform_ptr, waveform_ptr + waveform_elems); int32_t sample_length = frame_length * samples_per_frame_; if (sample_length < static_cast(audio->size())) { audio->resize(sample_length); @@ -717,13 +811,18 @@ class OfflineTtsKokoroModelAxcl::Impl { std::unique_ptr f0n_shared_model_; std::unique_ptr f0n_head_model_; std::unique_ptr decoder_front_model_; - std::unique_ptr vocoder_model_; + std::unique_ptr vocoder_core_model_; std::unique_ptr duration_predictor_; + std::unique_ptr vocoder_tail_; std::vector duration_input_names_; std::vector duration_input_names_ptr_; std::vector duration_output_names_; std::vector duration_output_names_ptr_; + std::vector vocoder_tail_input_names_; + std::vector vocoder_tail_input_names_ptr_; + std::vector vocoder_tail_output_names_; + std::vector vocoder_tail_output_names_ptr_; OfflineTtsKokoroModelMetaData meta_data_; std::vector style_dim_; @@ -737,6 +836,8 @@ class OfflineTtsKokoroModelAxcl::Impl { int32_t timbre_dim_ = 0; int32_t har_channels_ = 0; int32_t har_frames_ = 0; + int64_t vocoder_hidden_channels_ = 0; + int64_t vocoder_hidden_frames_ = 0; int32_t upsample_scale_ = 0; int32_t samples_per_frame_ = 0; }; diff --git a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc index 928f72e7c3..b37f45bb24 100644 --- a/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc +++ b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc @@ -78,6 +78,19 @@ static int32_t GetNumElements(const std::vector &shape) { return ans; } +static int64_t GetNumElements(const std::vector &shape) { + if (shape.empty()) { + return 0; + } + + int64_t ans = 1; + for (auto dim : shape) { + ans *= dim; + } + + return ans; +} + static std::string Basename(const std::string &path) { auto pos = path.find_last_of("/\\"); if (pos == std::string::npos) { @@ -158,6 +171,28 @@ static std::vector BuildPaddedTextMask(int32_t token_len, return mask; } +static std::vector GetOrtTensorShape(Ort::Session *sess, + const std::string &name, + bool is_input) { + Ort::AllocatorWithDefaultOptions allocator; + size_t count = is_input ? sess->GetInputCount() : sess->GetOutputCount(); + for (size_t i = 0; i != count; ++i) { + Ort::AllocatedStringPtr tensor_name = + is_input ? sess->GetInputNameAllocated(i, allocator) + : sess->GetOutputNameAllocated(i, allocator); + + if (name == tensor_name.get()) { + auto type_info = + is_input ? sess->GetInputTypeInfo(i) : sess->GetOutputTypeInfo(i); + return type_info.GetTensorTypeAndShapeInfo().GetShape(); + } + } + + SHERPA_ONNX_LOGE("Found no %s tensor with name: '%s'", + is_input ? "input" : "output", name.c_str()); + return {}; +} + struct ModelPaths { std::string encoder; std::string duration_predictor; @@ -165,7 +200,8 @@ struct ModelPaths { std::string f0n_shared; std::string f0n_head; std::string decoder_front; - std::string vocoder; + std::string vocoder_core; + std::string vocoder_tail; }; static ModelPaths ParseModelPaths(const std::string &model_str) { @@ -192,21 +228,24 @@ static ModelPaths ParseModelPaths(const std::string &model_str) { paths.f0n_head = f; } else if (name.find("decoder_front") != std::string::npos) { paths.decoder_front = f; - } else if (name.find("vocoder") != std::string::npos) { - paths.vocoder = f; + } else if (name.find("vocoder_tail") != std::string::npos) { + paths.vocoder_tail = f; + } else if (name.find("vocoder_core") != std::string::npos) { + paths.vocoder_core = f; } } if (paths.encoder.empty() || paths.duration_predictor.empty() || paths.text_encoder.empty() || paths.f0n_shared.empty() || paths.f0n_head.empty() || paths.decoder_front.empty() || - paths.vocoder.empty()) { + paths.vocoder_core.empty() || paths.vocoder_tail.empty()) { SHERPA_ONNX_LOGE( "For the new Axera Kokoro pipeline, --kokoro-model must contain these " - "7 files in any order: encoder_token_*.axmodel, " + "8 files in any order: encoder_token_*.axmodel, " "duration_predictor.onnx, text_encoder_token_*_frame_*.axmodel, " "f0n_shared_frame_*.axmodel, f0n_head_frame_*.axmodel, " - "decoder_front*.axmodel, vocoder*.axmodel. Given: %s", + "decoder_front*.axmodel, vocoder_core*.axmodel, " + "vocoder_tail*.onnx. Given: %s", model_str.c_str()); SHERPA_ONNX_EXIT(-1); } @@ -421,9 +460,17 @@ class OfflineTtsKokoroModelAxera::Impl { decoder_front_model_ = std::make_unique(decoder_front_buf.data(), decoder_front_buf.size()); - auto vocoder_buf = read_bytes(paths.vocoder); - vocoder_model_ = - std::make_unique(vocoder_buf.data(), vocoder_buf.size()); + auto vocoder_core_buf = read_bytes(paths.vocoder_core); + vocoder_core_model_ = std::make_unique(vocoder_core_buf.data(), + vocoder_core_buf.size()); + + auto vocoder_tail_buf = read_bytes(paths.vocoder_tail); + vocoder_tail_ = std::make_unique( + env_, vocoder_tail_buf.data(), vocoder_tail_buf.size(), sess_opts_); + GetInputNames(vocoder_tail_.get(), &vocoder_tail_input_names_, + &vocoder_tail_input_names_ptr_); + GetOutputNames(vocoder_tail_.get(), &vocoder_tail_output_names_, + &vocoder_tail_output_names_ptr_); auto duration_buf = read_bytes(paths.duration_predictor); duration_predictor_ = std::make_unique( @@ -449,14 +496,27 @@ class OfflineTtsKokoroModelAxera::Impl { auto f0n_head_output_shape = f0n_head_model_->TensorShape("F0_pred"); auto decoder_asr_shape = decoder_front_model_->TensorShape("asr"); auto decoder_timbre_shape = decoder_front_model_->TensorShape("timbre"); - auto vocoder_har_shape = vocoder_model_->TensorShape("har"); - auto vocoder_output_shape = vocoder_model_->TensorShape("waveform"); + auto vocoder_core_state_shape = + vocoder_core_model_->TensorShape("decoder_state"); + auto vocoder_core_timbre_shape = vocoder_core_model_->TensorShape("timbre"); + auto vocoder_core_har_shape = vocoder_core_model_->TensorShape("har"); + auto vocoder_core_hidden_shape = + vocoder_core_model_->TensorShape("vocoder_hidden"); + auto vocoder_tail_hidden_shape = + GetOrtTensorShape(vocoder_tail_.get(), "vocoder_hidden", true); + auto vocoder_tail_output_shape = + GetOrtTensorShape(vocoder_tail_.get(), "waveform", false); if (encoder_input_shape.size() != 2 || encoder_mask_shape.size() != 2 || text_encoder_input_shape.size() != 2 || text_encoder_aln_shape.size() != 3 || f0n_shared_input_shape.size() != 3 || f0n_head_ref_shape.size() != 2 || f0n_head_output_shape.size() != 2 || decoder_asr_shape.size() != 3 || - decoder_timbre_shape.size() != 2 || vocoder_har_shape.size() != 3) { + decoder_timbre_shape.size() != 2 || vocoder_core_state_shape.size() != 3 || + vocoder_core_timbre_shape.size() != 2 || + vocoder_core_har_shape.size() != 3 || + vocoder_core_hidden_shape.size() != 3 || + vocoder_tail_hidden_shape.size() != 3 || + vocoder_tail_output_shape.size() != 1) { SHERPA_ONNX_LOGE("Unexpected Kokoro Axera model shapes in: %s", config_.kokoro.model.c_str()); SHERPA_ONNX_EXIT(-1); @@ -467,9 +527,11 @@ class OfflineTtsKokoroModelAxera::Impl { pitch_bucket_ = f0n_head_output_shape[1]; ref_s_dim_ = f0n_head_ref_shape[1]; timbre_dim_ = decoder_timbre_shape[1]; - har_channels_ = vocoder_har_shape[1]; - har_frames_ = vocoder_har_shape[2]; - sample_bucket_ = GetNumElements(vocoder_output_shape); + har_channels_ = vocoder_core_har_shape[1]; + har_frames_ = vocoder_core_har_shape[2]; + vocoder_hidden_channels_ = vocoder_core_hidden_shape[1]; + vocoder_hidden_frames_ = vocoder_core_hidden_shape[2]; + sample_bucket_ = static_cast(GetNumElements(vocoder_tail_output_shape)); if (token_bucket_ != text_encoder_input_shape[1] || token_bucket_ != text_encoder_aln_shape[1] || @@ -486,6 +548,15 @@ class OfflineTtsKokoroModelAxera::Impl { SHERPA_ONNX_EXIT(-1); } + if (vocoder_core_state_shape != decoder_front_model_->TensorShape("decoder_state") || + vocoder_core_timbre_shape[1] != timbre_dim_ || + vocoder_tail_hidden_shape[1] != vocoder_hidden_channels_ || + vocoder_tail_hidden_shape[2] != vocoder_hidden_frames_) { + SHERPA_ONNX_LOGE("Kokoro Axera vocoder_core/tail shape mismatch in: %s", + config_.kokoro.model.c_str()); + SHERPA_ONNX_EXIT(-1); + } + if (har_channels_ != kStftNfft + 2) { SHERPA_ONNX_LOGE("Expected har channels %d. Given: %d", kStftNfft + 2, har_channels_); @@ -812,16 +883,39 @@ class OfflineTtsKokoroModelAxera::Impl { decoder_front_model_->GetOutputTensorData("decoder_state"); std::vector har = BuildHar(f0_pred); - if (!vocoder_model_->SetInputTensorData("decoder_state", decoder_state.data(), - decoder_state.size()) || - !vocoder_model_->SetInputTensorData("timbre", timbre.data(), - timbre.size()) || - !vocoder_model_->SetInputTensorData("har", har.data(), har.size()) || - !vocoder_model_->Run()) { + if (!vocoder_core_model_->SetInputTensorData("decoder_state", + decoder_state.data(), + decoder_state.size()) || + !vocoder_core_model_->SetInputTensorData("timbre", timbre.data(), + timbre.size()) || + !vocoder_core_model_->SetInputTensorData("har", har.data(), har.size()) || + !vocoder_core_model_->Run()) { + return false; + } + + std::vector vocoder_hidden = + vocoder_core_model_->GetOutputTensorData("vocoder_hidden"); + std::array vocoder_hidden_shape = { + 1, vocoder_hidden_channels_, vocoder_hidden_frames_}; + Ort::Value vocoder_hidden_tensor = Ort::Value::CreateTensor( + memory_info, vocoder_hidden.data(), vocoder_hidden.size(), + vocoder_hidden_shape.data(), vocoder_hidden_shape.size()); + std::array vocoder_tail_inputs = { + std::move(vocoder_hidden_tensor)}; + + auto vocoder_tail_outputs = vocoder_tail_->Run( + {}, vocoder_tail_input_names_ptr_.data(), vocoder_tail_inputs.data(), + vocoder_tail_inputs.size(), vocoder_tail_output_names_ptr_.data(), + vocoder_tail_output_names_ptr_.size()); + if (vocoder_tail_outputs.empty()) { + SHERPA_ONNX_LOGE("Unexpected empty vocoder_tail outputs"); return false; } - *audio = vocoder_model_->GetOutputTensorData("waveform"); + const float *waveform_ptr = vocoder_tail_outputs[0].GetTensorData(); + int32_t waveform_elems = static_cast( + vocoder_tail_outputs[0].GetTensorTypeAndShapeInfo().GetElementCount()); + *audio = std::vector(waveform_ptr, waveform_ptr + waveform_elems); int32_t sample_length = frame_length * samples_per_frame_; if (sample_length < static_cast(audio->size())) { audio->resize(sample_length); @@ -844,13 +938,18 @@ class OfflineTtsKokoroModelAxera::Impl { std::unique_ptr f0n_shared_model_; std::unique_ptr f0n_head_model_; std::unique_ptr decoder_front_model_; - std::unique_ptr vocoder_model_; + std::unique_ptr vocoder_core_model_; std::unique_ptr duration_predictor_; + std::unique_ptr vocoder_tail_; std::vector duration_input_names_; std::vector duration_input_names_ptr_; std::vector duration_output_names_; std::vector duration_output_names_ptr_; + std::vector vocoder_tail_input_names_; + std::vector vocoder_tail_input_names_ptr_; + std::vector vocoder_tail_output_names_; + std::vector vocoder_tail_output_names_ptr_; OfflineTtsKokoroModelMetaData meta_data_; std::vector style_dim_; @@ -864,6 +963,8 @@ class OfflineTtsKokoroModelAxera::Impl { int32_t timbre_dim_ = 0; int32_t har_channels_ = 0; int32_t har_frames_ = 0; + int64_t vocoder_hidden_channels_ = 0; + int64_t vocoder_hidden_frames_ = 0; int32_t upsample_scale_ = 0; int32_t samples_per_frame_ = 0; };