From 68dfaf07de6ff43f24b35dcca4e3dafbfbde83b8 Mon Sep 17 00:00:00 2001
From: jslok <jslok930@gmail.com>
Date: Sun, 13 Sep 2026 22:44:14 -0700
Subject: [PATCH] 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 <memory>
+#include <utility>
+#include <vector>
 
 #if defined(ANDROID)
 #include <tflite/c/c_api.h>
+#include <tflite/delegates/gpu/delegate.h>
+#include <tflite/delegates/nnapi/nnapi_delegate_c_api.h>
 #elif defined(__APPLE__)
 #include <TensorFlowLiteC/TensorFlowLiteC.h>
+#if FAST_TFLITE_ENABLE_CORE_ML
+#include <TensorFlowLiteCCoreML/TensorFlowLiteCCoreML.h>
+#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<TfLiteDelegate, DelegateDeleter>;
+
 std::shared_ptr<HybridTfliteModelSpec>
 HybridTfliteModule::createModel(const std::shared_ptr<ArrayBuffer>& modelData,
                                 const std::vector<TensorflowModelDelegate>& delegates) {
@@ -50,20 +88,39 @@ HybridTfliteModule::createModel(const std::shared_ptr<ArrayBuffer>& modelData,
 
   // Add all hardware accelerated delegates (e.g. GPU, NPU, ...)
   // if any. The default CPU delegate will always be available.
+  std::vector<TensorflowModelDelegate> effectiveDelegates;
+  std::vector<OwnedDelegate> 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<TfLiteInterpreter> 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<ArrayBuffer> to keep model data bytes alive
-  return std::make_shared<HybridTfliteModel>(interpreter, modelData, delegates);
+  // Wrap in HybridTfliteModel — stores shared_ptr<ArrayBuffer> to keep model data bytes alive.
+  // Only the delegates that were actually registered are reported via `getDelegates()`.
+  return std::make_shared<HybridTfliteModel>(interpreter, modelData, effectiveDelegates);
 }
 
 } // namespace margelo::nitro::tflite