From 68dfaf07de6ff43f24b35dcca4e3dafbfbde83b8 Mon Sep 17 00:00:00 2001 From: jslok Date: Sun, 13 Sep 2026 22:44:14 -0700 Subject: [PATCH 1/2] fix: skip unavailable delegates and free delegates on model destruction Two related delegate-lifecycle fixes in createModel: 1. Skip null delegates: delegate factories can legitimately return nullptr (e.g. TfLiteCoreMlDelegateCreate on devices without a Neural Engine when enabled_devices is ANE-only). Registering that nullptr with TfLiteInterpreterOptionsAddDelegate crashes/corrupts the interpreter. Now a null delegate is skipped so the model falls back to CPU, and getDelegates() reports only the delegates that were actually registered. 2. Free delegates: TFLite's C API does not transfer delegate ownership to the interpreter - the caller must delete delegates itself after the interpreter is destroyed. They were never freed, so every model destruction leaked the delegate's compiled kernels / driver contexts (GPU: TfLiteGpuDelegateV2Delete, NNAPI: TfLiteNnapiDelegateDelete, CoreML: TfLiteCoreMlDelegateDelete). Each delegate is now held in a unique_ptr with its own delete function; ownership moves into the interpreter's shared_ptr deleter, so delegates are freed right after TfLiteInterpreterDelete (they must outlive the interpreter) and on every failure path in createModel via normal unwinding. --- cpp/HybridTfliteModule.cpp | 67 +++++++++++++++++++++++++++++++++++--- 1 file changed, 62 insertions(+), 5 deletions(-) diff --git a/cpp/HybridTfliteModule.cpp b/cpp/HybridTfliteModule.cpp index 91102012..631ee879 100644 --- a/cpp/HybridTfliteModule.cpp +++ b/cpp/HybridTfliteModule.cpp @@ -2,11 +2,18 @@ #include "TfliteHelpers.hpp" #include +#include +#include #if defined(ANDROID) #include +#include +#include #elif defined(__APPLE__) #include +#if FAST_TFLITE_ENABLE_CORE_ML +#include +#endif #else #error "Invalid Platform!" #endif @@ -32,6 +39,37 @@ TfLiteDelegate* getDelegate(TensorflowModelDelegate delegateType) { "\"!"); } +/** + * TFLite's C API does not transfer delegate ownership to the interpreter: the + * caller must keep a delegate alive for the interpreter's lifetime and free it + * afterwards with the delegate's own delete function. + */ +struct DelegateDeleter { + TensorflowModelDelegate delegateType; + + void operator()(TfLiteDelegate* delegate) const { + switch (delegateType) { +#if defined(__APPLE__) && FAST_TFLITE_ENABLE_CORE_ML + case TensorflowModelDelegate::CORE_ML: + TfLiteCoreMlDelegateDelete(delegate); + return; +#endif +#if defined(ANDROID) + case TensorflowModelDelegate::ANDROID_GPU: + TfLiteGpuDelegateV2Delete(delegate); + return; + case TensorflowModelDelegate::NNAPI: + TfLiteNnapiDelegateDelete(delegate); + return; +#endif + default: + // getDelegate() throws for every other type on this platform. + return; + } + } +}; +using OwnedDelegate = std::unique_ptr; + std::shared_ptr HybridTfliteModule::createModel(const std::shared_ptr& modelData, const std::vector& delegates) { @@ -50,20 +88,39 @@ HybridTfliteModule::createModel(const std::shared_ptr& modelData, // Add all hardware accelerated delegates (e.g. GPU, NPU, ...) // if any. The default CPU delegate will always be available. + std::vector effectiveDelegates; + std::vector ownedDelegates; + effectiveDelegates.reserve(delegates.size()); + ownedDelegates.reserve(delegates.size()); for (const TensorflowModelDelegate& delegateType : delegates) { - TfLiteDelegate* delegate = getDelegate(delegateType); - TfLiteInterpreterOptionsAddDelegate(options.get(), delegate); + OwnedDelegate delegate(getDelegate(delegateType), DelegateDeleter{delegateType}); + if (delegate == nullptr) { + // e.g. CoreML on devices without a Neural Engine — fall back to CPU + // instead of registering a null delegate with the interpreter. + continue; + } + TfLiteInterpreterOptionsAddDelegate(options.get(), delegate.get()); + effectiveDelegates.push_back(delegateType); + ownedDelegates.push_back(std::move(delegate)); } TfLiteInterpreter* rawInterpreter = TfLiteInterpreterCreate(model.get(), options.get()); if (rawInterpreter == nullptr) { + // `ownedDelegates` frees the delegates on unwind. throw std::runtime_error("Failed to create TFLite interpreter!"); } + // The delegates travel with the interpreter and are freed right after it, + // so they can never be deleted while the interpreter still uses them. const std::shared_ptr interpreter( - rawInterpreter, [modelData](TfLiteInterpreter* value) { TfLiteInterpreterDelete(value); }); + rawInterpreter, + [modelData, ownedDelegates = std::move(ownedDelegates)](TfLiteInterpreter* value) mutable { + TfLiteInterpreterDelete(value); + ownedDelegates.clear(); + }); - // Wrap in HybridTfliteModel — stores shared_ptr to keep model data bytes alive - return std::make_shared(interpreter, modelData, delegates); + // Wrap in HybridTfliteModel — stores shared_ptr to keep model data bytes alive. + // Only the delegates that were actually registered are reported via `getDelegates()`. + return std::make_shared(interpreter, modelData, effectiveDelegates); } } // namespace margelo::nitro::tflite From bc2c9d452059409b3770f9877078d5082a5150b9 Mon Sep 17 00:00:00 2001 From: jslok Date: Sun, 13 Sep 2026 22:45:15 -0700 Subject: [PATCH 2/2] feat: implement dispose() for deterministic native resource release HybridTfliteModel inherits Nitro's default no-op dispose(), so JS calling model.dispose() frees nothing - the interpreter, delegates and model buffer only go away when GC drops the last reference. Worklet runtimes (frame processors) may not GC for a long time, especially while the app is backgrounded, so multi-hundred-MB models and their GPU contexts stay resident with no way to release them deterministically. This implements a real dispose(): - Resets the interpreter shared_ptr immediately. The model is its only owner, so the interpreter deleter runs right away: TfLiteInterpreterDelete, then the delegates. Our references to the model bytes and the cached output buffers are dropped too. - Thread-safe via a lifecycle mutex: dispose() blocks until an in-flight inference on another thread completes - freeing the interpreter under a running TfLiteInterpreterInvoke would be a native crash. The mutex is uncontended in normal operation (~ns per lock vs ~ms per inference). - run() re-checks disposal on the async thread, since dispose() may land between the caller-thread input copy and the async invoke. - All post-dispose calls (runSync/run/getInputs/getOutputs) throw a catchable JS error ('TFLite: Model was disposed!') instead of crashing. - Idempotent; the destructor stays defaulted since a disposed model holds nothing. --- cpp/HybridTfliteModel.cpp | 37 ++++++++++++++++++++++++++++++++++--- cpp/HybridTfliteModel.hpp | 25 +++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 3 deletions(-) diff --git a/cpp/HybridTfliteModel.cpp b/cpp/HybridTfliteModel.cpp index 020b60fa..439de90e 100644 --- a/cpp/HybridTfliteModel.cpp +++ b/cpp/HybridTfliteModel.cpp @@ -28,11 +28,28 @@ HybridTfliteModel::HybridTfliteModel(std::shared_ptr interpre } } +void HybridTfliteModel::dispose() { + // The lock waits out an in-flight inference on another thread - freeing the + // interpreter under a running TfLiteInterpreterInvoke would be a native crash. + std::lock_guard lock(_lifecycleMutex); + if (_interpreter == nullptr) { + return; // already disposed + } + // We hold the only reference to the interpreter, so this runs its deleter + // now: TfLiteInterpreterDelete, then the delegates. The model bytes are + // released once nobody else (e.g. JS) references them either. + _interpreter.reset(); + _modelData.reset(); + _outputBuffers.clear(); +} + std::vector HybridTfliteModel::getDelegates() { return _delegates; } std::vector HybridTfliteModel::getInputs() { + std::lock_guard lock(_lifecycleMutex); + throwIfDisposed(); int count = TfLiteInterpreterGetInputTensorCount(_interpreter.get()); std::vector tensors; tensors.reserve(count); @@ -55,6 +72,8 @@ std::vector HybridTfliteModel::getInputs() { } std::vector HybridTfliteModel::getOutputs() { + std::lock_guard lock(_lifecycleMutex); + throwIfDisposed(); int count = TfLiteInterpreterGetOutputTensorCount(_interpreter.get()); std::vector tensors; tensors.reserve(count); @@ -148,6 +167,10 @@ void HybridTfliteModel::invoke() { std::vector> HybridTfliteModel::runSync(const std::vector>& input) { + // Held for the whole inference so dispose() can never free the interpreter + // mid-invoke. The disposed-throw is a catchable JS error on any runtime. + std::lock_guard lock(_lifecycleMutex); + throwIfDisposed(); copyInputBuffers(input); invoke(); return copyOutputBuffers(); @@ -155,12 +178,20 @@ HybridTfliteModel::runSync(const std::vector>& inpu std::shared_ptr>>> HybridTfliteModel::run(const std::vector>& input) { - // Copy input buffers on caller (JS) thread first — input ArrayBuffers are - // non-owning JS buffers that may be GC'd if we access them async. - copyInputBuffers(input); + { + // Copy input buffers on caller (JS) thread first — input ArrayBuffers are + // non-owning JS buffers that may be GC'd if we access them async. + std::lock_guard lock(_lifecycleMutex); + throwIfDisposed(); + copyInputBuffers(input); + } std::shared_ptr sharedThis = shared_cast(); return Promise>>::async( [sharedThis]() -> std::vector> { + // Re-acquire on the async thread: dispose() may have landed between + // the input copy above and this lambda running. + std::lock_guard lock(sharedThis->_lifecycleMutex); + sharedThis->throwIfDisposed(); sharedThis->invoke(); return sharedThis->copyOutputBuffers(); }); diff --git a/cpp/HybridTfliteModel.hpp b/cpp/HybridTfliteModel.hpp index d0a137ca..b9b04969 100644 --- a/cpp/HybridTfliteModel.hpp +++ b/cpp/HybridTfliteModel.hpp @@ -2,6 +2,8 @@ #include "HybridTfliteModelSpec.hpp" #include +#include +#include #include #include @@ -22,6 +24,17 @@ class HybridTfliteModel : public HybridTfliteModelSpec { std::vector delegates); ~HybridTfliteModel() override = default; + /** + * Free the interpreter, its delegates and our reference to the model bytes + * NOW, without waiting for every runtime's GC to drop its reference (worklet + * runtimes may not GC for a long time, especially while backgrounded). + * Thread-safe: blocks until an in-flight inference on another thread has + * completed. Every later runSync/run/getInputs/getOutputs call throws a + * catchable JS error. Idempotent. Called by Nitro when JS invokes + * `model.dispose()`. + */ + void dispose() override; + // Properties (from HybridTfliteModelSpec) std::vector getDelegates() override; std::vector getInputs() override; @@ -39,11 +52,23 @@ class HybridTfliteModel : public HybridTfliteModelSpec { std::vector> copyOutputBuffers(); std::shared_ptr getOutputBufferForTensor(const TfLiteTensor* tensor); + // Caller must hold _lifecycleMutex. A disposed model has released its + // interpreter; std::runtime_error surfaces as a catchable JS error. + void throwIfDisposed() const { + if (_interpreter == nullptr) { + throw std::runtime_error("TFLite: Model was disposed!"); + } + } + private: std::shared_ptr _interpreter; std::vector _delegates; std::shared_ptr _modelData; std::unordered_map> _outputBuffers; + + // Serializes inference against dispose(). Uncontended in normal operation + // (one lock per inference, ~ns vs ~ms inference cost). + std::mutex _lifecycleMutex; }; } // namespace margelo::nitro::tflite