From cac462c48660367bd53e72da076c6c72a665eb45 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Sat, 2 Sep 2023 09:54:10 -0700
Subject: [PATCH] Add allow_custom_ops to model_util.convert_to_tflite and
 enable custom ops for face stylizer.

PiperOrigin-RevId: 562212965
---
 .../model_maker/python/core/utils/model_util.py     | 13 ++++++++++++-
 .../python/vision/face_stylizer/face_stylizer.py    |  6 ++----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/mediapipe/model_maker/python/core/utils/model_util.py b/mediapipe/model_maker/python/core/utils/model_util.py
index fd11c60b2..58262ed9b 100644
--- a/mediapipe/model_maker/python/core/utils/model_util.py
+++ b/mediapipe/model_maker/python/core/utils/model_util.py
@@ -119,6 +119,7 @@ def convert_to_tflite_from_file(
         tf.lite.OpsSet.TFLITE_BUILTINS,
     ),
     preprocess: Optional[Callable[..., Any]] = None,
+    allow_custom_ops: bool = False,
 ) -> bytearray:
   """Converts the input Keras model to TFLite format.
 
@@ -129,6 +130,8 @@ def convert_to_tflite_from_file(
     preprocess: A callable to preprocess the representative dataset for
       quantization. The callable takes three arguments in order: feature, label,
       and is_training.
+    allow_custom_ops: A boolean flag to enable custom ops in model convsion.
+      Default to False.
 
   Returns:
     bytearray of TFLite model
@@ -140,6 +143,7 @@ def convert_to_tflite_from_file(
         converter, preprocess=preprocess
     )
 
+  converter.allow_custom_ops = allow_custom_ops
   converter.target_spec.supported_ops = supported_ops
   tflite_model = converter.convert()
   return tflite_model
@@ -152,6 +156,7 @@ def convert_to_tflite(
         tf.lite.OpsSet.TFLITE_BUILTINS,
     ),
     preprocess: Optional[Callable[..., Any]] = None,
+    allow_custom_ops: bool = False,
 ) -> bytearray:
   """Converts the input Keras model to TFLite format.
 
@@ -162,6 +167,8 @@ def convert_to_tflite(
     preprocess: A callable to preprocess the representative dataset for
       quantization. The callable takes three arguments in order: feature, label,
       and is_training.
+    allow_custom_ops: A boolean flag to enable custom ops in model conversion.
+      Default to False.
 
   Returns:
     bytearray of TFLite model
@@ -174,7 +181,11 @@ def convert_to_tflite(
         save_format='tf',
     )
     return convert_to_tflite_from_file(
-        save_path, quantization_config, supported_ops, preprocess
+        save_path,
+        quantization_config,
+        supported_ops,
+        preprocess,
+        allow_custom_ops,
     )
 
 
diff --git a/mediapipe/model_maker/python/vision/face_stylizer/face_stylizer.py b/mediapipe/model_maker/python/vision/face_stylizer/face_stylizer.py
index e635d2b32..b23a6e498 100644
--- a/mediapipe/model_maker/python/vision/face_stylizer/face_stylizer.py
+++ b/mediapipe/model_maker/python/vision/face_stylizer/face_stylizer.py
@@ -274,11 +274,9 @@ class FaceStylizer(object):
     face_stylizer_model_buffer = model_util.convert_to_tflite(
         model=model,
         quantization_config=None,
-        supported_ops=(
-            tf.lite.OpsSet.TFLITE_BUILTINS,
-            tf.lite.OpsSet.SELECT_TF_OPS,
-        ),
+        supported_ops=(tf.lite.OpsSet.TFLITE_BUILTINS,),
         preprocess=self._preprocessor,
+        allow_custom_ops=True,
     )
 
     face_aligner_task_file_path = constants.FACE_ALIGNER_TASK_FILES.get_path()