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 diff --git a/sherpa-onnx/csrc/CMakeLists.txt b/sherpa-onnx/csrc/CMakeLists.txt index 1a34d1c1b0..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) @@ -223,6 +228,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-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/axcl/offline-tts-kokoro-model-axcl.cc b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc new file mode 100644 index 0000000000..f85d0527f2 --- /dev/null +++ b/sherpa-onnx/csrc/axcl/offline-tts-kokoro-model-axcl.cc @@ -0,0 +1,876 @@ +// 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 +#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 "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" +#include "sherpa-onnx/csrc/text-utils.h" + +namespace sherpa_onnx { + +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 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) { + return path; + } + + 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; +} + +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 ans; +} + +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; +} + +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; + std::string text_encoder; + std::string f0n_shared; + std::string f0n_head; + std::string decoder_front; + std::string vocoder_core; + std::string vocoder_tail; +}; + +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_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_core.empty() || paths.vocoder_tail.empty()) { + SHERPA_ONNX_LOGE( + "For the new Axcl Kokoro pipeline, --kokoro-model must contain these " + "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_core*.axmodel, " + "vocoder_tail*.onnx. 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, [](const std::string &filename) { return ReadFile(filename); }); + } + + template + Impl(Manager *mgr, const OfflineTtsModelConfig &config) + : config_(config), + env_(ORT_LOGGING_LEVEL_ERROR), + sess_opts_(GetSessionOptions(config)), + allocator_{} { + Init(config, [mgr](const std::string &filename) { return ReadFile(mgr, filename); }); + } + + const OfflineTtsKokoroModelMetaData &GetMetaData() const { + return meta_data_; + } + + 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.0f / 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; + if (!RunModels(input_ids, ref_s, speed, &audio)) { + SHERPA_ONNX_LOGE("Run Kokoro Axcl static pipeline failed"); + return {}; + } + + return audio; + } + + private: + 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); + } + + 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()); + + 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_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( + 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()); + } + + 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_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_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); + } + + 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_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] || + 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); + } + + 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); + } + + 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_); + SHERPA_ONNX_EXIT(-1); + } + + if (sample_bucket_ % pitch_bucket_ != 0 || sample_bucket_ % frame_bucket_ != 0) { + SHERPA_ONNX_LOGE( + "Unexpected Kokoro Axcl bucket relation. sample=%d pitch=%d frame=%d", + sample_bucket_, pitch_bucket_, frame_bucket_); + SHERPA_ONNX_EXIT(-1); + } + + 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); + } + + 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); + } + } + + 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); + } + + 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); + } + + 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) 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); + } + + 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; + } + + 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 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; + } + + 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; + } + + 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]; + } + + sine_merge[sample_index] = std::tanh(merged); + } + + 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::array window{}; + std::array, kStftNfft / 2 + 1> cos_table{}; + std::array, kStftNfft / 2 + 1> sin_table{}; + + for (int32_t n = 0; n != kStftNfft; ++n) { + window[n] = 0.5f - 0.5f * std::cos(kTwoPi * n / static_cast(kStftNfft)); + } + + 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 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]; + } + + 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; + } + } + + 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]); + } + + 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 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 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; + } + + 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; + } + + 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; + } + + std::vector asr = text_encoder_model_->GetOutputTensorData("asr"); + + 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 (!f0n_shared_model_->SetInputTensorData("en", padded_en.data(), + padded_en.size()) || + !f0n_shared_model_->Run()) { + return false; + } + + 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; + } + + 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; + } + + std::vector decoder_state = + decoder_front_model_->GetOutputTensorData("decoder_state"); + std::vector har = BuildHar(f0_pred); + + 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; + } + + 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); + } + + return true; + } + + private: + mutable std::mutex mutex_; + + OfflineTtsModelConfig config_; + Ort::Env env_; + Ort::SessionOptions sess_opts_; + Ort::AllocatorWithDefaultOptions allocator_; + + 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_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_; + std::vector styles_; + + 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; + int64_t vocoder_hidden_channels_ = 0; + int64_t vocoder_hidden_frames_ = 0; + int32_t upsample_scale_ = 0; + int32_t samples_per_frame_ = 0; +}; + +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/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..b37f45bb24 --- /dev/null +++ b/sherpa-onnx/csrc/axera/offline-tts-kokoro-model-axera.cc @@ -0,0 +1,1003 @@ +// 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 +#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 "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" +#include "sherpa-onnx/csrc/text-utils.h" + +namespace sherpa_onnx { + +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 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) { + return path; + } + + 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; +} + +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 ans; +} + +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; +} + +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; + std::string text_encoder; + std::string f0n_shared; + std::string f0n_head; + std::string decoder_front; + std::string vocoder_core; + std::string vocoder_tail; +}; + +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_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_core.empty() || paths.vocoder_tail.empty()) { + SHERPA_ONNX_LOGE( + "For the new Axera Kokoro pipeline, --kokoro-model must contain these " + "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_core*.axmodel, " + "vocoder_tail*.onnx. Given: %s", + model_str.c_str()); + SHERPA_ONNX_EXIT(-1); + } + + return paths; +} + +} // namespace + +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() { + int 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_; +}; + +class OfflineTtsKokoroModelAxera::Impl { + public: + explicit Impl(const OfflineTtsModelConfig &config) + : config_(config), + env_(ORT_LOGGING_LEVEL_ERROR), + sess_opts_(GetSessionOptions(config)), + allocator_{} { + Init(config, [](const std::string &filename) { return ReadFile(filename); }); + } + + template + Impl(Manager *mgr, const OfflineTtsModelConfig &config) + : config_(config), + env_(ORT_LOGGING_LEVEL_ERROR), + sess_opts_(GetSessionOptions(config)), + allocator_{} { + Init(config, [mgr](const std::string &filename) { return ReadFile(mgr, filename); }); + } + + const OfflineTtsKokoroModelMetaData &GetMetaData() const { + return meta_data_; + } + + 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.0f / 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; + if (!RunModels(input_ids, ref_s, speed, &audio)) { + SHERPA_ONNX_LOGE("Run Kokoro Axera static pipeline failed"); + return {}; + } + + return audio; + } + + private: + 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); + } + + 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()); + + 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_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( + 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()); + } + + 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_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_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); + } + + 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_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] || + 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); + } + + 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); + } + + 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_); + SHERPA_ONNX_EXIT(-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); + } + + 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); + } + + 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); + } + } + + 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); + } + + 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); + } + + 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) 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); + } + + 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; + } + + 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 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; + } + + 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; + } + + 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]; + } + + sine_merge[sample_index] = std::tanh(merged); + } + + 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::array window{}; + std::array, kStftNfft / 2 + 1> cos_table{}; + std::array, kStftNfft / 2 + 1> sin_table{}; + + for (int32_t n = 0; n != kStftNfft; ++n) { + window[n] = 0.5f - 0.5f * std::cos(kTwoPi * n / static_cast(kStftNfft)); + } + + 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 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]; + } + + 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; + } + } + + 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 Axera 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]); + } + + 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 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 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; + } + + 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; + } + + 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; + } + + std::vector asr = text_encoder_model_->GetOutputTensorData("asr"); + + 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 (!f0n_shared_model_->SetInputTensorData("en", padded_en.data(), + padded_en.size()) || + !f0n_shared_model_->Run()) { + return false; + } + + 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; + } + + 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; + } + + std::vector decoder_state = + decoder_front_model_->GetOutputTensorData("decoder_state"); + std::vector har = BuildHar(f0_pred); + + 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; + } + + 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); + } + + return true; + } + + private: + mutable std::mutex mutex_; + AxEngineGuard ax_engine_guard_; + + OfflineTtsModelConfig config_; + Ort::Env env_; + Ort::SessionOptions sess_opts_; + Ort::AllocatorWithDefaultOptions allocator_; + + 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_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_; + std::vector styles_; + + 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; + int64_t vocoder_hidden_channels_ = 0; + int64_t vocoder_hidden_frames_ = 0; + int32_t upsample_scale_ = 0; + int32_t samples_per_frame_ = 0; +}; + +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 7720a49113..1fcdd4b3c8 100644 --- a/sherpa-onnx/csrc/offline-tts-impl.cc +++ b/sherpa-onnx/csrc/offline-tts-impl.cc @@ -24,6 +24,14 @@ #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 + namespace sherpa_onnx { std::vector OfflineTtsImpl::AddBlank(const std::vector &x, @@ -40,6 +48,51 @@ 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< + OfflineTtsKokoroAxclImpl>(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()) { @@ -65,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()) { 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()) {