react-native-nitro-onnx 0.1.0

package/cpp/SpeakerManager.hpp

@@ -0,0 +1,43 @@
+// ------------------------------------------------------------------------------
+// SpeakerManager.hpp
+// ------------------------------------------------------------------------------
+#pragma once
+
+#include "SpeakerEngine.hpp"
+
+#include <NitroModules/Promise.hpp>
+#include <NitroModules/HybridObject.hpp>
+#include <HybridSpeakerManagerSpec.hpp>
+
+#include <memory>
+
+namespace margelo::nitro::onnx::speech {
+
+class SpeakerManager : public HybridSpeakerManagerSpec {
+ public:
+  static constexpr auto TAG = "SpeakerManager";
+
+  SpeakerManager(std::shared_ptr<ThreadPool> threadPool, std::string cacheDir);
+  ~SpeakerManager() override;
+
+  std::shared_ptr<Promise<void>> load(const SpeakerEmbeddingConfig& config) override;
+  bool isLoaded() override;
+  std::shared_ptr<Promise<std::shared_ptr<ArrayBuffer>>> computeEmbedding(
+      const std::shared_ptr<ArrayBuffer>& samples) override;
+  std::shared_ptr<Promise<RegisteredSpeaker>> registerSpeaker(
+      const std::string& id,
+      const std::string& name,
+      const std::shared_ptr<ArrayBuffer>& embedding) override;
+  std::shared_ptr<Promise<RegisteredSpeaker>> registerSpeakerFromFile(
+      const std::string& id,
+      const std::string& name,
+      const std::string& path) override;
+  std::shared_ptr<Promise<std::vector<RegisteredSpeaker>>> listSpeakers() override;
+  std::shared_ptr<Promise<void>> removeSpeaker(const std::string& id) override;
+  std::shared_ptr<Promise<void>> unload() override;
+
+ private:
+  SpeakerEngine engine_;
+};
+
+}  // namespace margelo::nitro::onnx::speech

package/cpp/StreamingAsr.cpp

@@ -0,0 +1,123 @@
+// ------------------------------------------------------------------------------
+// StreamingAsr.cpp
+// ------------------------------------------------------------------------------
+#include "StreamingAsr.hpp"
+
+#include "AudioUtils.hpp"
+
+#include <optional>
+
+namespace margelo::nitro::onnx::speech {
+
+namespace {
+
+AsrResult toAsrResult(const AsrEngineResult& native) {
+  std::vector<double> timestamps;
+  timestamps.reserve(native.timestamps.size());
+  for (float t : native.timestamps) {
+    timestamps.push_back(static_cast<double>(t));
+  }
+  return AsrResult(
+      native.text,
+      std::optional<double>(static_cast<double>(native.score)),
+      std::optional<double>(static_cast<double>(native.startMs)),
+      std::optional<double>(static_cast<double>(native.endMs)),
+      timestamps.empty() ? std::nullopt : std::optional<std::vector<double>>(timestamps),
+      native.json.empty() ? std::nullopt : std::optional<std::string>(native.json));
+}
+
+}  // namespace
+
+StreamingAsr::StreamingAsr(std::shared_ptr<ThreadPool> threadPool)
+    : HybridObject(TAG), engine_(std::move(threadPool)) {}
+
+StreamingAsr::~StreamingAsr() {
+  engine_.unload();
+}
+
+std::shared_ptr<Promise<void>> StreamingAsr::load(const AsrModelConfig& config) {
+  return Promise<void>::async([this, config]() {
+    AsrEngineConfig native;
+    native.type = config.type;
+    native.modelDir = config.modelDir;
+    native.tokensPath = config.tokensPath;
+    native.encoder = config.encoder.value_or("");
+    native.decoder = config.decoder.value_or("");
+    native.joiner = config.joiner.value_or("");
+    native.numThreads = static_cast<int32_t>(config.numThreads.value_or(2));
+    native.decodingMethod = config.decodingMethod.value_or("greedy_search");
+    native.maxActivePaths = static_cast<int32_t>(config.maxActivePaths.value_or(4));
+    engine_.load(native, shared_cast<StreamingAsr>());
+  });
+}
+
+bool StreamingAsr::isLoaded() {
+  return engine_.isLoaded();
+}
+
+std::shared_ptr<Promise<void>> StreamingAsr::acceptWaveform(
+    const std::shared_ptr<ArrayBuffer>& samples) {
+  return Promise<void>::async([this, samples]() {
+    auto floatSamples = bytesToFloatVector(samples->data(), samples->size());
+    engine_.acceptWaveform(floatSamples);
+  });
+}
+
+std::shared_ptr<Promise<AsrResult>> StreamingAsr::finalize() {
+  return Promise<AsrResult>::async([this]() { return toAsrResult(engine_.finalize()); });
+}
+
+std::shared_ptr<Promise<void>> StreamingAsr::reset() {
+  return Promise<void>::async([this]() { engine_.reset(); });
+}
+
+std::shared_ptr<Promise<void>> StreamingAsr::unload() {
+  return Promise<void>::async([this]() { engine_.unload(); });
+}
+
+std::optional<std::function<void(const AsrResult& /* result */)>> StreamingAsr::getOnPartialResult() {
+  return onPartialResult_;
+}
+
+void StreamingAsr::setOnPartialResult(
+    const std::optional<std::function<void(const AsrResult& /* result */)>>& onPartialResult) {
+  onPartialResult_ = onPartialResult;
+}
+
+std::optional<std::function<void(const AsrResult& /* result */)>> StreamingAsr::getOnFinalResult() {
+  return onFinalResult_;
+}
+
+void StreamingAsr::setOnFinalResult(
+    const std::optional<std::function<void(const AsrResult& /* result */)>>& onFinalResult) {
+  onFinalResult_ = onFinalResult;
+}
+
+std::optional<std::function<void(const std::string& /* error */)>> StreamingAsr::getOnError() {
+  return onError_;
+}
+
+void StreamingAsr::setOnError(
+    const std::optional<std::function<void(const std::string& /* error */)>>& onError) {
+  onError_ = onError;
+}
+
+void StreamingAsr::onPartialResult(const AsrEngineResult& result) {
+  if (onPartialResult_.has_value()) {
+    onPartialResult_.value()(toAsrResult(result));
+  }
+}
+
+void StreamingAsr::onFinalResult(const AsrEngineResult& result) {
+  if (onFinalResult_.has_value()) {
+    onFinalResult_.value()(toAsrResult(result));
+  }
+}
+
+void StreamingAsr::onError(const std::string& error) {
+  if (onError_.has_value()) {
+    onError_.value()(error);
+  }
+}
+
+}  // namespace margelo::nitro::onnx::speech

package/cpp/StreamingAsr.hpp

@@ -0,0 +1,55 @@
+// ------------------------------------------------------------------------------
+// StreamingAsr.hpp
+// ------------------------------------------------------------------------------
+#pragma once
+
+#include "AsrEngine.hpp"
+
+#include <NitroModules/Promise.hpp>
+#include <NitroModules/HybridObject.hpp>
+#include <HybridStreamingAsrSpec.hpp>
+
+#include <functional>
+#include <memory>
+#include <optional>
+
+namespace margelo::nitro::onnx::speech {
+
+class StreamingAsr : public HybridStreamingAsrSpec,
+                         public StreamingAsrListener {
+ public:
+  static constexpr auto TAG = "StreamingAsr";
+
+  explicit StreamingAsr(std::shared_ptr<ThreadPool> threadPool);
+  ~StreamingAsr() override;
+
+  std::shared_ptr<Promise<void>> load(const AsrModelConfig& config) override;
+  bool isLoaded() override;
+  std::shared_ptr<Promise<void>> acceptWaveform(const std::shared_ptr<ArrayBuffer>& samples) override;
+  std::shared_ptr<Promise<AsrResult>> finalize() override;
+  std::shared_ptr<Promise<void>> reset() override;
+  std::shared_ptr<Promise<void>> unload() override;
+
+  // Nitro event callback getters/setters.
+  std::optional<std::function<void(const AsrResult& /* result */)>> getOnPartialResult() override;
+  void setOnPartialResult(
+      const std::optional<std::function<void(const AsrResult& /* result */)>>& onPartialResult) override;
+  std::optional<std::function<void(const AsrResult& /* result */)>> getOnFinalResult() override;
+  void setOnFinalResult(
+      const std::optional<std::function<void(const AsrResult& /* result */)>>& onFinalResult) override;
+  std::optional<std::function<void(const std::string& /* error */)>> getOnError() override;
+  void setOnError(
+      const std::optional<std::function<void(const std::string& /* error */)>>& onError) override;
+
+  void onPartialResult(const AsrEngineResult& result) override;
+  void onFinalResult(const AsrEngineResult& result) override;
+  void onError(const std::string& error) override;
+
+ private:
+  StreamingAsrEngine engine_;
+  std::optional<std::function<void(const AsrResult&)>> onPartialResult_;
+  std::optional<std::function<void(const AsrResult&)>> onFinalResult_;
+  std::optional<std::function<void(const std::string&)>> onError_;
+};
+
+}  // namespace margelo::nitro::onnx::speech

package/cpp/ThreadPool.cpp

@@ -0,0 +1,41 @@
+// ------------------------------------------------------------------------------
+// ThreadPool.cpp
+// ------------------------------------------------------------------------------
+#include "ThreadPool.hpp"
+
+namespace margelo::nitro::onnx::speech {
+
+ThreadPool::ThreadPool(size_t threadCount) {
+  for (size_t i = 0; i < threadCount; ++i) {
+    workers_.emplace_back([this]() {
+      for (;;) {
+        std::function<void()> task;
+        {
+          std::unique_lock<std::mutex> lock(queueMutex_);
+          condition_.wait(lock, [this]() { return stop_ || !tasks_.empty(); });
+          if (stop_ && tasks_.empty()) {
+            return;
+          }
+          task = std::move(tasks_.front());
+          tasks_.pop();
+        }
+        task();
+      }
+    });
+  }
+}
+
+ThreadPool::~ThreadPool() {
+  {
+    std::unique_lock<std::mutex> lock(queueMutex_);
+    stop_ = true;
+  }
+  condition_.notify_all();
+  for (std::thread& worker : workers_) {
+    if (worker.joinable()) {
+      worker.join();
+    }
+  }
+}
+
+}  // namespace margelo::nitro::onnx::speech

package/cpp/ThreadPool.hpp

@@ -0,0 +1,62 @@
+// ------------------------------------------------------------------------------
+// ThreadPool.hpp
+// A fixed-size thread pool used to run all heavy native inference tasks.
+// Keeping inference off the JS thread is a core requirement of this module.
+// ------------------------------------------------------------------------------
+#pragma once
+
+#include <atomic>
+#include <condition_variable>
+#include <functional>
+#include <future>
+#include <mutex>
+#include <queue>
+#include <thread>
+#include <vector>
+
+namespace margelo::nitro::onnx::speech {
+
+/**
+ * Minimal thread pool for background inference work.
+ * Tasks are enqueued as std::function<void()> and executed by worker threads.
+ */
+class ThreadPool final {
+ public:
+  explicit ThreadPool(size_t threadCount = std::thread::hardware_concurrency());
+  ~ThreadPool();
+
+  ThreadPool(const ThreadPool&) = delete;
+  ThreadPool& operator=(const ThreadPool&) = delete;
+
+  /**
+   * Schedule a task on the pool and obtain a future for the result.
+   * @tparam F Callable type.
+   * @tparam Args Argument types.
+   * @return std::future for the callable result.
+   */
+  template <typename F, typename... Args>
+  auto enqueue(F&& f, Args&&... args) -> std::future<std::invoke_result_t<F, Args...>> {
+    using ReturnType = std::invoke_result_t<F, Args...>;
+    auto task = std::make_shared<std::packaged_task<ReturnType()>>(
+        std::bind(std::forward<F>(f), std::forward<Args>(args)...));
+    std::future<ReturnType> result = task->get_future();
+    {
+      std::unique_lock<std::mutex> lock(queueMutex_);
+      if (stop_) {
+        throw std::runtime_error("Cannot enqueue on stopped ThreadPool");
+      }
+      tasks_.emplace([task]() { (*task)(); });
+    }
+    condition_.notify_one();
+    return result;
+  }
+
+ private:
+  std::vector<std::thread> workers_;
+  std::queue<std::function<void()>> tasks_;
+  std::mutex queueMutex_;
+  std::condition_variable condition_;
+  std::atomic<bool> stop_{false};
+};
+
+}  // namespace margelo::nitro::onnx::speech

package/cpp/Tts.cpp

@@ -0,0 +1,132 @@
+// ------------------------------------------------------------------------------
+// TtsImpl.cpp
+// ------------------------------------------------------------------------------
+#include "Tts.hpp"
+
+#include "AudioUtils.hpp"
+
+#include <NitroModules/ArrayBuffer.hpp>
+#include <fstream>
+#include <stdexcept>
+
+namespace margelo::nitro::onnx::speech {
+
+namespace {
+
+TtsResult toTtsResult(const TtsEngineResult& native) {
+  std::vector<uint8_t> bytes = floatVectorToBytes(native.samples);
+  return TtsResult(
+      ArrayBuffer::move(std::move(bytes)),
+      static_cast<double>(native.sampleRate),
+      native.durationMs);
+}
+
+}  // namespace
+
+Tts::Tts(std::shared_ptr<ThreadPool> threadPool)
+    : HybridObject(TAG), engine_(std::move(threadPool)) {}
+
+Tts::~Tts() {
+  engine_.unload();
+}
+
+std::shared_ptr<Promise<void>> Tts::load(const TtsModelConfig& config) {
+  return Promise<void>::async([this, config]() {
+    TtsEngineConfig native;
+    native.type = static_cast<TtsModelType>(config.type);
+    native.modelDir = config.modelDir;
+    native.model = config.model.value_or("");
+    native.acousticModel = config.acousticModel.value_or("");
+    native.vocoder = config.vocoder.value_or("");
+    native.tokens = config.tokens.value_or("");
+    native.lexicon = config.lexicon.value_or("");
+    native.voices = config.voices.value_or("");
+    native.espeakNgData = config.espeakNgData.value_or("");
+    native.dictDir = config.dictDir.value_or("");
+    native.lmMain = config.lmMain.value_or("");
+    native.lmFlow = config.lmFlow.value_or("");
+    native.textConditioner = config.textConditioner.value_or("");
+    native.pocketEncoder = config.pocketEncoder.value_or("");
+    native.pocketDecoder = config.pocketDecoder.value_or("");
+    native.vocabJson = config.vocabJson.value_or("");
+    native.tokenScoresJson = config.tokenScoresJson.value_or("");
+    native.zipvoiceEncoder = config.zipvoiceEncoder.value_or("");
+    native.zipvoiceDecoder = config.zipvoiceDecoder.value_or("");
+    native.config = config.config.value_or("");
+    native.numThreads = static_cast<int32_t>(config.numThreads.value_or(2));
+    native.outputSampleRate = static_cast<int32_t>(config.outputSampleRate.value_or(16000.0));
+    native.speakerId = static_cast<int32_t>(config.speakerId.value_or(0.0));
+    native.speed = static_cast<float>(config.speed.value_or(1.0));
+    engine_.load(native);
+  });
+}
+
+bool Tts::isLoaded() {
+  return engine_.isLoaded();
+}
+
+std::shared_ptr<Promise<TtsResult>> Tts::synthesize(const std::string& text, std::optional<double> speed) {
+  return Promise<TtsResult>::async([this, text, speed]() {
+    float spd = speed.has_value() ? static_cast<float>(speed.value()) : -1.0f;
+    return toTtsResult(engine_.synthesize(text, -1, spd));
+  });
+}
+
+std::shared_ptr<Promise<TtsResult>> Tts::synthesizeWithSpeaker(
+    const std::string& text,
+    const std::string& speakerId,
+    std::optional<double> speed) {
+  return Promise<TtsResult>::async([this, text, speakerId, speed]() {
+    int32_t sid = std::stoi(speakerId);
+    float spd = speed.has_value() ? static_cast<float>(speed.value()) : -1.0f;
+    return toTtsResult(engine_.synthesize(text, sid, spd));
+  });
+}
+
+std::shared_ptr<Promise<void>> Tts::saveWav(const TtsResult& result, const std::string& path) {
+  return Promise<void>::async([result, path]() {
+    const auto sampleRate = static_cast<uint32_t>(result.sampleRate);
+    const uint16_t numChannels = 1;
+    const uint16_t bitsPerSample = 32;
+    const uint16_t audioFormat = 3; // IEEE float
+
+    const auto* sampleData = reinterpret_cast<const uint8_t*>(result.samples->data());
+    const uint32_t dataByteCount = static_cast<uint32_t>(result.samples->size());
+
+    const uint32_t byteRate = sampleRate * numChannels * bitsPerSample / 8;
+    const uint16_t blockAlign = numChannels * bitsPerSample / 8;
+    const uint32_t chunkSize = 36 + dataByteCount;
+
+    std::ofstream file(path, std::ios::binary);
+    if (!file) {
+      throw std::runtime_error("Cannot open file for writing: " + path);
+    }
+
+    // RIFF header
+    file.write("RIFF", 4);
+    file.write(reinterpret_cast<const char*>(&chunkSize), 4);
+    file.write("WAVE", 4);
+
+    // fmt chunk
+    file.write("fmt ", 4);
+    uint32_t fmtSize = 16;
+    file.write(reinterpret_cast<const char*>(&fmtSize), 4);
+    file.write(reinterpret_cast<const char*>(&audioFormat), 2);
+    file.write(reinterpret_cast<const char*>(&numChannels), 2);
+    file.write(reinterpret_cast<const char*>(&sampleRate), 4);
+    file.write(reinterpret_cast<const char*>(&byteRate), 4);
+    file.write(reinterpret_cast<const char*>(&blockAlign), 2);
+    file.write(reinterpret_cast<const char*>(&bitsPerSample), 2);
+
+    // data chunk
+    file.write("data", 4);
+    file.write(reinterpret_cast<const char*>(&dataByteCount), 4);
+    file.write(reinterpret_cast<const char*>(sampleData), dataByteCount);
+  });
+}
+
+std::shared_ptr<Promise<void>> Tts::unload() {
+  return Promise<void>::async([this]() { engine_.unload(); });
+}
+
+}  // namespace margelo::nitro::onnx::speech

package/cpp/Tts.hpp

@@ -0,0 +1,38 @@
+// ------------------------------------------------------------------------------
+// TtsImpl.hpp
+// ------------------------------------------------------------------------------
+#pragma once
+
+#include "TtsEngine.hpp"
+
+#include <NitroModules/Promise.hpp>
+#include <NitroModules/HybridObject.hpp>
+#include <HybridTtsSpec.hpp>
+
+#include <memory>
+#include <optional>
+
+namespace margelo::nitro::onnx::speech {
+
+class Tts : public HybridTtsSpec {
+ public:
+  static constexpr auto TAG = "Tts";
+
+  explicit Tts(std::shared_ptr<ThreadPool> threadPool);
+  ~Tts() override;
+
+  std::shared_ptr<Promise<void>> load(const TtsModelConfig& config) override;
+  bool isLoaded() override;
+  std::shared_ptr<Promise<TtsResult>> synthesize(const std::string& text, std::optional<double> speed) override;
+  std::shared_ptr<Promise<TtsResult>> synthesizeWithSpeaker(
+      const std::string& text,
+      const std::string& speakerId,
+      std::optional<double> speed) override;
+  std::shared_ptr<Promise<void>> saveWav(const TtsResult& result, const std::string& path) override;
+  std::shared_ptr<Promise<void>> unload() override;
+
+ private:
+  TtsEngine engine_;
+};
+
+}  // namespace margelo::nitro::onnx::speech