Update tensorflow_models usage of tf.lite.interpreter to run ai-edge-…

…litert.interpreter PiperOrigin-RevId: 682467705
tensorflow · Oct 4, 2024 · 6eef699 · 6eef699
1 parent 8f60a93
commit 6eef699
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/requirements.txt b/requirements.txt
@@ -13,3 +13,4 @@ termcolor >= 1.1.0
 trimesh >= 2.37.22
 # Required by trimesh.
 networkx
+ai-edge-litert >= 1.0.1
diff --git a/tensorflow_graphics/util/test_case.py b/tensorflow_graphics/util/test_case.py
@@ -31,6 +31,10 @@
 import tensorflow as tf
 
 from tensorflow_graphics.util import tfg_flags
+# pylint: disable=g-direct-tensorflow-import
+from ai-edge-litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
+
 
 FLAGS = flags.FLAGS
 
@@ -364,7 +368,7 @@ def assert_tf_lite_convertible(self,
             sess, in_tensors, out_tensors)
         tflite_model = converter.convert()
         # Load TFLite model and allocate tensors.
-        interpreter = tf.lite.Interpreter(model_content=tflite_model)
+        interpreter = tfl_interpreter.Interpreter(model_content=tflite_model)
         interpreter.allocate_tensors()
         # If no test inputs provided then randomly generate inputs.
         if test_inputs is None: