Add internal kwarg to set intermediate saved model dir

PiperOrigin-RevId: 706748161

Author: chunnienc

ai_edge_torch/_convert/conversion.py

@@ -78,7 +78,8 @@ def convert_signatures(
     *,
     strict_export: Union[Literal["auto"], bool] = True,
     quant_config: Optional[qcfg.QuantConfig] = None,
-    _tfl_converter_flags: Optional[dict[str, Any]],
+    _tfl_converter_flags: Optional[dict[str, Any]] = None,
+    _saved_model_dir: Optional[str] = None,
 ) -> model.TfLiteModel:
   """Converts a list of `signature.Signature`s and embeds them into one `model.TfLiteModel`.
 
@@ -93,6 +94,8 @@ def convert_signatures(
       quant_config: User-defined quantization method and scheme of the model.
       _tfl_converter_flags: A nested dictionary allowing setting flags for the
         underlying tflite converter.
+      _saved_model_dir: Directory for the intermediate saved model. If not
+        specified, a random temporary directory would be used.
 
   Returns:
     The converted `model.TfLiteModel` object.
@@ -140,6 +143,7 @@ def export(*args, **kwargs):
       signatures,
       quant_config=quant_config,
       _tfl_converter_flags=_tfl_converter_flags,
+      _saved_model_dir=_saved_model_dir,
   )
 
   return model.TfLiteModel(tflite_model)

ai_edge_torch/_convert/converter.py

@@ -106,6 +106,7 @@ def convert(
       quant_config: Optional[qcfg.QuantConfig] = None,
       dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
       _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
+      _saved_model_dir: Optional[str] = None,
   ) -> model.TfLiteModel:
     """Finalizes the conversion and produces an edge model.
 
@@ -139,6 +140,8 @@ def convert(
         of this function and so needs to be treated as such. Please do not rely
         on this parameter except for local debugging as this can be removed in a
         future release.
+      _saved_model_dir: Directory for the intermediate saved model. If not
+        specified, a random temporary directory would be used.
 
     Returns:
       The converted edge model.
@@ -171,6 +174,7 @@ def convert(
         strict_export=strict_export,
         quant_config=quant_config,
         _tfl_converter_flags=_ai_edge_converter_flags,
+        _saved_model_dir=_saved_model_dir,
     )
 
 
@@ -216,6 +220,7 @@ def convert(
     quant_config: Optional[qcfg.QuantConfig] = None,
     dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
     _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
+    _saved_model_dir: Optional[str] = None,
 ) -> model.TfLiteModel:
   """Converts a PyTorch model to an edge model with a default signature.
 
@@ -240,6 +245,8 @@ def convert(
       this function and so needs to be treated as such. Please do not rely on
       this parameter except for local debugging as this can be removed in a
       future release.
+    _saved_model_dir: Directory for the intermediate saved model. If not
+      specified, a random temporary directory would be used.
 
   Returns:
     The converted edge model.
@@ -259,4 +266,5 @@ def convert(
       quant_config=quant_config,
       dynamic_shapes=dynamic_shapes,
       _ai_edge_converter_flags=_ai_edge_converter_flags,
+      _saved_model_dir=_saved_model_dir,
   )

ai_edge_torch/lowertools/_shim.py

@@ -50,13 +50,16 @@ def exported_programs_to_tflite(
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: Optional[dict[str, Any]] = None,
+    _saved_model_dir: Optional[str] = None
 ):
   """Converts a list of ExportedProgram to a TFLite model.
 
   Args:
     exported_programs: A list of ExportedProgram.
     signatures: A list of Signature.
     quant_config: A QuantConfig.
+    _saved_model_dir: Directory for the intermediate saved model. If not
+      specified, a random temporary directory would be used.
     _tfl_converter_flags: A dict of flags for TFLiteConverter.
 
   Returns:
@@ -79,4 +82,5 @@ def exported_programs_to_tflite(
       signatures,
       quant_config=quant_config,
       _tfl_converter_flags=_tfl_converter_flags,
+      _saved_model_dir=_saved_model_dir,
   )

ai_edge_torch/lowertools/odml_torch_utils.py

@@ -138,6 +138,7 @@ def merged_bundle_to_tfl_model(
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: dict = {},
+    _saved_model_dir: Optional[str] = None,
 ):
   tf_state_dict = merged_bundle.bundles[0].state_dict
 
@@ -173,6 +174,9 @@ def merged_bundle_to_tfl_model(
   # We need to temporarily save since TFLite's from_concrete_functions does not
   # allow providing names for each of the concrete functions.
   with tempfile.TemporaryDirectory() as temp_dir_path:
+    if _saved_model_dir is not None:
+      temp_dir_path = _saved_model_dir
+
     tf.saved_model.save(
         tf_module,
         temp_dir_path,

ai_edge_torch/lowertools/torch_xla_utils.py

@@ -192,6 +192,7 @@ def merged_bundle_to_tfl_model(
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: dict = {},
+    _saved_model_dir: Optional[str] = None,
 ) -> None:
   """Converts a StableHLOGraphModule to a tflite model.
 
@@ -200,6 +201,8 @@ def merged_bundle_to_tfl_model(
     signatures: List of signatures from which names of the signatures is
     extracted.
     quant_config: User-defined quantization method and scheme of the model.
+    _saved_model_dir: Directory for the intermediate saved model. If not
+      specified, a random temporary directory would be used.
     _tfl_converter_flags: A nested dictionary allowing setting flags for the
     underlying tflite converter.
   """
@@ -246,6 +249,9 @@ def merged_bundle_to_tfl_model(
   # We need to temporarily save since TFLite's from_concrete_functions does not
   # allow providing names for each of the concrete functions.
   with tempfile.TemporaryDirectory() as temp_dir_path:
+    if _saved_model_dir is not None:
+      temp_dir_path = _saved_model_dir
+
     tf.saved_model.save(
         tf_module,
         temp_dir_path,