Move tensorflow lite python calls to ai-edge-litert.

PiperOrigin-RevId: 690670455
google-ai-edge · Oct 28, 2024 · fbc4999 · fbc4999
1 parent 51d1f3b
commit fbc4999
Show file tree

Hide file tree

Showing 6 changed files with 4 additions and 15 deletions.
diff --git a/mediapipe/model_maker/python/core/utils/BUILD b/mediapipe/model_maker/python/core/utils/BUILD
@@ -44,7 +44,6 @@ py_library(
     deps = [
         ":quantization",
         "//mediapipe/model_maker/python/core/data:dataset",
-        "@model_maker_pip_deps_ai_edge_litert_nightly//:pkg",
         "@model_maker_pip_deps_numpy//:pkg",
         "@model_maker_pip_deps_tensorflow//:pkg",
     ],
@@ -173,7 +172,6 @@ py_test(
         ":quantization",
         ":test_util",
         "@model_maker_pip_deps_absl_py//:pkg",
-        "@model_maker_pip_deps_ai_edge_litert_nightly//:pkg",
         "@model_maker_pip_deps_tensorflow//:pkg",
     ],
 )
diff --git a/mediapipe/model_maker/python/core/utils/model_util.py b/mediapipe/model_maker/python/core/utils/model_util.py
@@ -28,7 +28,6 @@
 
 from mediapipe.model_maker.python.core.data import dataset
 from mediapipe.model_maker.python.core.utils import quantization
-from ai_edge_litert import interpreter as tfl_interpreter
 
 DEFAULT_SCALE, DEFAULT_ZERO_POINT = 0, 0
 ESTIMITED_STEPS_PER_EPOCH = 1000
@@ -274,7 +273,7 @@ def __init__(self, tflite_model: bytearray):
     Args:
       tflite_model: A valid flatbuffer representing the TFLite model.
     """
-    self.interpreter = tfl_interpreter.Interpreter(model_content=tflite_model)
+    self.interpreter = tf.lite.Interpreter(model_content=tflite_model)
     self.interpreter.allocate_tensors()
     self.input_details = self.interpreter.get_input_details()
     self.output_details = self.interpreter.get_output_details()

diff --git a/mediapipe/model_maker/python/core/utils/quantization_test.py b/mediapipe/model_maker/python/core/utils/quantization_test.py
@@ -17,7 +17,6 @@
 
 from mediapipe.model_maker.python.core.utils import quantization
 from mediapipe.model_maker.python.core.utils import test_util
-from ai_edge_litert import interpreter as tfl_interpreter
 
 
 class QuantizationTest(tf.test.TestCase, parameterized.TestCase):
@@ -60,7 +59,7 @@ def test_set_converter_with_quantization_from_int8_config(self):
     self.assertEqual(config.supported_ops,
                      [tf.lite.OpsSet.TFLITE_BUILTINS_INT8])
     tflite_model = converter.convert()
-    interpreter = tfl_interpreter.Interpreter(model_content=tflite_model)
+    interpreter = tf.lite.Interpreter(model_content=tflite_model)
     self.assertEqual(interpreter.get_input_details()[0]['dtype'], tf.uint8)
     self.assertEqual(interpreter.get_output_details()[0]['dtype'], tf.uint8)
 
@@ -83,7 +82,7 @@ def test_set_converter_with_quantization_from_float16_config(self):
     converter = config.set_converter_with_quantization(converter=converter)
     self.assertEqual(config.supported_types, [tf.float16])
     tflite_model = converter.convert()
-    interpreter = tfl_interpreter.Interpreter(model_content=tflite_model)
+    interpreter = tf.lite.Interpreter(model_content=tflite_model)
     # The input and output are expected to be set to float32 by default.
     self.assertEqual(interpreter.get_input_details()[0]['dtype'], tf.float32)
     self.assertEqual(interpreter.get_output_details()[0]['dtype'], tf.float32)

diff --git a/mediapipe/model_maker/requirements.txt b/mediapipe/model_maker/requirements.txt
@@ -1,5 +1,4 @@
 absl-py
-ai-edge-litert-nightly
 mediapipe>=0.10.0
 numpy<2
 opencv-python

diff --git a/mediapipe/model_maker/requirements_bazel.txt b/mediapipe/model_maker/requirements_bazel.txt
@@ -1,5 +1,4 @@
 absl-py
-ai-edge-litert-nightly
 numpy<2
 opencv-python
 setuptools==70.3.0 # needed due to https://github.com/pypa/setuptools/issues/4487

diff --git a/mediapipe/model_maker/requirements_lock.txt b/mediapipe/model_maker/requirements_lock.txt
@@ -15,8 +15,6 @@ absl-py==1.4.0
     #   tensorflow-metadata
     #   tensorflow-model-optimization
     #   tf-slim
-ai-edge-litert-nightly==1.0.1.dev20241022
-    # via -r mediapipe/opensource_only/model_maker_requirements_bazel.txt
 array-record==0.5.1
     # via tensorflow-datasets
 astunparse==1.6.3
@@ -50,9 +48,7 @@ etils[array-types,enp,epath,epy,etqdm,etree]==1.5.2
     #   array-record
     #   tensorflow-datasets
 flatbuffers==24.3.25
-    # via
-    #   ai-edge-litert-nightly
-    #   tensorflow
+    # via tensorflow
 fonttools==4.54.1
     # via matplotlib
 fsspec==2024.9.0
@@ -124,7 +120,6 @@ ml-dtypes==0.3.2
 numpy==1.26.4
     # via
     #   -r mediapipe/opensource_only/model_maker_requirements_bazel.txt
-    #   ai-edge-litert-nightly
     #   contourpy
     #   etils
     #   h5py