ai-edge-litert-nightly 1.4.0.dev20250813__cp39-cp39-manylinux_2_27_aarch64.whl → 1.4.0.dev20250814__cp39-cp39-manylinux_2_27_aarch64.whl

ai_edge_litert/__init__.py

@@ -1 +1 @@
-__version__ = "1.4.0.dev20250813"
+__version__ = "1.4.0.dev20250814"

{ai_edge_litert_nightly-1.4.0.dev20250813.dist-info → ai_edge_litert_nightly-1.4.0.dev20250814.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-edge-litert-nightly
-Version: 1.4.0.dev20250813
+Version: 1.4.0.dev20250814
 Summary: LiteRT is for mobile and embedded devices.
 Home-page: https://www.tensorflow.org/lite/
 Author: Google AI Edge Authors

{ai_edge_litert_nightly-1.4.0.dev20250813.dist-info → ai_edge_litert_nightly-1.4.0.dev20250814.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ai_edge_litert/__init__.py,sha256=IldzDEKlTmxzLLoGNY0HlSynO81y6L2E-oKhiyTvgZc,34
+ai_edge_litert/__init__.py,sha256=_8wnWQPu-Jzh5AhwdogIQzRrnYwgDBb0b51yQ0pL8ng,34
 ai_edge_litert/_pywrap_analyzer_wrapper.so,sha256=qhOo4FG4sHlk0LRBUnHeXCZKhHGKgrzmA0WjmIFpPe8,6464
 ai_edge_litert/_pywrap_litert_compiled_model_wrapper.so,sha256=RRsmEAo9KuLT8lx0o8pIrS7zbwtYhHNLsvTx5hTyA-0,6488
 ai_edge_litert/_pywrap_litert_tensor_buffer_wrapper.so,sha256=-3w_TOE2Fnav4vc2ZN3aXtqcDhJP8RXpImo7AmUrpBk,6488
@@ -17,8 +17,7 @@ ai_edge_litert/field_mask_pb2.py,sha256=YVi1IasKjEQAlcIwM6HMgmW0wT6TMAXjiaiysrEY
 ai_edge_litert/format_converter_wrapper_pybind11.so,sha256=tIhiUYnHh5ZKIDTuMKKuOMwqxd6MjrPN7EMYV7oNO-M,6512
 ai_edge_litert/hardware_accelerator.py,sha256=FgouirT1Cs80mqvMvLbO9O7wiPBnQaWDB5vjFpJhtJo,800
 ai_edge_litert/interpreter.py,sha256=4u9FfiaYkDxf_IxCjzeLWyTyxwG184uVxqg0fbj6tYg,40650
-ai_edge_litert/libLiteRtRuntimeCApi.so,sha256=bH-yyyoIM8RPfSIuI_iCuHIa-cywj0saiHn7eA1NuXM,8920576
-ai_edge_litert/libpywrap_litert_common.so,sha256=SQTEkbQxmuayam5cVbqzOUgBjObS1Bg0uHJCxsXMfo8,14716272
+ai_edge_litert/libpywrap_litert_common.so,sha256=91z0KyAIZZH8ipaDL_MjV987Hodwhc68HoHpQnmcDOE,15317624
 ai_edge_litert/metrics_interface.py,sha256=dVu6SmbnQUntPgE5o6BxHVMyemwli-7F6tDfVMGrlYI,1542
 ai_edge_litert/metrics_portable.py,sha256=KKvR9ZOe8j2ZeBtDo_6gWJ8kENKoOawPK3LPkevnZa8,2039
 ai_edge_litert/model_runtime_info_pb2.py,sha256=E93kYJtWnsChrdegZJbKzeFpplssBGEURTdOOfjtOxg,6370
@@ -32,31 +31,7 @@ ai_edge_litert/tensor_buffer.py,sha256=3Xy7kd_aR499QrfO0YITvr_z-aU1eONga41KiDBUe
 ai_edge_litert/timestamp_pb2.py,sha256=K7Gs_qOn2XAMZTNUUrJ1XcGf4OmYhdt858I_zDtZYzg,1793
 ai_edge_litert/type_pb2.py,sha256=x8rTulmlk9FxtRJal7BnUsaPbn2TcLlu8D7AVoVPjMA,5065
 ai_edge_litert/wrappers_pb2.py,sha256=4hQAZRGeaE5DyY6YQ7VfrqozPa_d_LCOBEYjxfQNlKs,2955
-ai_edge_litert/aot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_edge_litert/aot/aot_compile.py,sha256=i0c6owaMBT7N2VUUSk8KrP35XX0R9-c5XzTHQQRBIcU,5501
-ai_edge_litert/aot/prepare_for_npu.py,sha256=bghtIcLXw5P6u_lHOh-oBv4u-jCoUi-kI3wF3wA1Mow,5494
-ai_edge_litert/aot/ai_pack/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_edge_litert/aot/ai_pack/export_lib.py,sha256=KVzfU2zZ0c6hx0dC_nLMPYUr8NROacknUH6e-MFMWfE,10237
-ai_edge_litert/aot/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_edge_litert/aot/core/apply_plugin.py,sha256=DAWj1I9k94WCkwophgz5Nrj8A0kSO2v8X3nqN_1KJFk,4554
-ai_edge_litert/aot/core/common.py,sha256=5imHTXV2abVwd5RJqFD0iw6WlLZ73njOPtCUccYtrWM,3103
-ai_edge_litert/aot/core/components.py,sha256=DcdEDEiA6U0HT5Tt10aw_j_M1I-XtRiXTwE5hGBMEiI,2415
-ai_edge_litert/aot/core/mlir_transforms.py,sha256=cW1DWZX1FDEI87Ihf3CfJ31lQmfZEBLLdptQCsK4sjM,1289
-ai_edge_litert/aot/core/tflxx_util.py,sha256=fyg_K1-nVaEx207sZ88m9TOp3VcjJWpSFknv9_S8cF0,1020
-ai_edge_litert/aot/core/types.py,sha256=Bvv27bUQPGQSZON2bQZ8BZTxoxLDvQki6oqtymp6fx4,10547
-ai_edge_litert/aot/vendors/__init__.py,sha256=K11L1REFCE0umYEUwmTxhijRfVbIytqOych2QgP4b78,864
-ai_edge_litert/aot/vendors/fallback_backend.py,sha256=LvvnXyR-EANTzh9DOYwSGstoYYcsQxtxNyc5vb9P7qE,3396
-ai_edge_litert/aot/vendors/import_vendor.py,sha256=HNDkx9vdmDGtX1SupfChKdW4UtCA8Xb6xC_S7xdBKkU,3469
-ai_edge_litert/aot/vendors/example/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_edge_litert/aot/vendors/example/example_backend.py,sha256=dgwTivJQzL1z3M_w4ixmOe_681E6cZW0bupgRPM_zOM,4182
-ai_edge_litert/aot/vendors/mediatek/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_edge_litert/aot/vendors/mediatek/mediatek_backend.py,sha256=uFSciDUollsKciJooxn5Fo7FWdAidSG0Q0BRBxE8TQI,6074
-ai_edge_litert/aot/vendors/mediatek/target.py,sha256=MvwdgybSHmVYuN1txgzmDUvO5HYubx9TLzOi1TXkJx4,2214
-ai_edge_litert/aot/vendors/qualcomm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_edge_litert/aot/vendors/qualcomm/qualcomm_backend.py,sha256=121jVb4u1slOQ8FSf1vJBNiq1OPIMeNB9cEQrrWByJk,4864
-ai_edge_litert/aot/vendors/qualcomm/target.py,sha256=BesYeOdRQILJY2GRVzenakTWQU-fCTqs9VPTTLAA8mw,1796
-ai_edge_litert/tools/apply_plugin_main,sha256=43hZ-dPuf-gqxNqsE_rBWTC_ZllZcmAtTB0L_29msOg,2561952
-ai_edge_litert_nightly-1.4.0.dev20250813.dist-info/METADATA,sha256=bkiv0hJ0iTLL5N-NIQsyGKNwAhlBZsv_cDpL6Kxq_yU,1911
-ai_edge_litert_nightly-1.4.0.dev20250813.dist-info/WHEEL,sha256=2V3Z1XExrBUQtUy2R-D_eV91jXaDJRE0crbndTl2Q00,112
-ai_edge_litert_nightly-1.4.0.dev20250813.dist-info/top_level.txt,sha256=WcDZgG99n0a0xDS_ipL8ZWy956g1v5fVyR3FH96VDT0,15
-ai_edge_litert_nightly-1.4.0.dev20250813.dist-info/RECORD,,
+ai_edge_litert_nightly-1.4.0.dev20250814.dist-info/METADATA,sha256=uUlRPU3OhyxzwH9YEEuY1K8ltKCBP4j4HCL9umv5zQ8,1911
+ai_edge_litert_nightly-1.4.0.dev20250814.dist-info/WHEEL,sha256=2V3Z1XExrBUQtUy2R-D_eV91jXaDJRE0crbndTl2Q00,112
+ai_edge_litert_nightly-1.4.0.dev20250814.dist-info/top_level.txt,sha256=WcDZgG99n0a0xDS_ipL8ZWy956g1v5fVyR3FH96VDT0,15
+ai_edge_litert_nightly-1.4.0.dev20250814.dist-info/RECORD,,

ai_edge_litert/aot/ai_pack/export_lib.py

@@ -1,281 +0,0 @@
-# Copyright 2025 Google LLC.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Utility functions for exporting models to AI pack format."""
-
-import itertools
-import os
-import pathlib
-from typing import cast
-
-from ai_edge_litert.aot.core import common
-from ai_edge_litert.aot.core import types
-from ai_edge_litert.aot.vendors import fallback_backend
-from ai_edge_litert.aot.vendors.mediatek import mediatek_backend
-from ai_edge_litert.aot.vendors.mediatek import target as mtk_target
-from ai_edge_litert.aot.vendors.qualcomm import qualcomm_backend
-from ai_edge_litert.aot.vendors.qualcomm import target as qnn_target
-
-# TODO: b/407453529 - Add unittests.
-
-
-_DEVICE_TARGETING_CONFIGURATION = """<config:device-targeting-config
-    xmlns:config="http://schemas.android.com/apk/config">
-{device_groups}
-</config:device-targeting-config>"""
-
-_DEVICE_GROUP_TEMPLATE = """    <config:device-group name="{device_group_name}">
-{device_selectors}
-    </config:device-group>"""
-
-_DEVICE_SELECTOR_TEMPLATE = """        <config:device-selector>
-            <config:system-on-chip manufacturer="{soc_man}" model="{soc_model}"/>
-        </config:device-selector>"""
-
-
-def _is_mobile_device_backend(backend: types.Backend):
-  target = backend.target
-  if backend.id() == qualcomm_backend.QualcommBackend.id():
-    target = cast(qnn_target.Target, target)
-    # Non Android QNN targets.
-    if target.soc_model in (
-        qnn_target.SocModel.SA8255,
-        qnn_target.SocModel.SA8295,
-    ):
-      return False
-  return True
-
-
-def _export_model_files_to_ai_pack(
-    compiled_models: types.CompilationResult,
-    ai_pack_dir: pathlib.Path,
-    ai_pack_name: str,
-    litert_model_name: str,
-    *,
-    separate_mtk_ai_pack: bool = True,
-):
-  """Exports the model tflite files to the AI pack directory structure.
-
-  Args:
-    compiled_models: The compiled models to export.
-    ai_pack_dir: The directory to export the AI pack to.
-    ai_pack_name: The name of the AI pack.
-    litert_model_name: The name of the model in the litert format.
-    separate_mtk_ai_pack: Whether to separate the MTK AI pack. If True, the main
-      AI pack will use the fallback model for MTK targets. The MTK AI pack will
-      contain all MTK models, and empty directories for non-MTK targets.
-  """
-  fallback_model = None
-  for backend, model in compiled_models.models_with_backend:
-    if backend.target_id == fallback_backend.FallbackBackend.id():
-      fallback_model = model
-  assert fallback_model is not None, 'Fallback model is required.'
-
-  model_export_dir = ai_pack_dir / ai_pack_name / 'src/main/assets'
-  os.makedirs(model_export_dir, exist_ok=True)
-  for backend, model in compiled_models.models_with_backend:
-    if not _is_mobile_device_backend(backend):
-      continue
-    target_id = backend.target_id
-    backend_id = backend.id()
-    if backend_id == fallback_backend.FallbackBackend.id():
-      target_id = 'other'
-    elif backend_id == mediatek_backend.MediaTekBackend.id():
-      target_id = backend.target_id.replace(
-          mtk_target.SocManufacturer.MEDIATEK, 'Mediatek'
-      )
-    group_name = 'model#group_' + target_id
-    export_dir = model_export_dir / group_name
-    os.makedirs(export_dir, exist_ok=True)
-    model_export_path = export_dir / (litert_model_name + common.DOT_TFLITE)
-    if (
-        separate_mtk_ai_pack
-        and backend_id == mediatek_backend.MediaTekBackend.id()
-    ):
-      # Use the fallback model for MTK targets in main AI pack.
-      model_to_export = fallback_model
-    else:
-      model_to_export = model
-    if not model_to_export.in_memory:
-      model_to_export.load()
-    model_to_export.save(model_export_path, export_only=True)
-
-  if separate_mtk_ai_pack:
-    _export_model_files_to_mtk_ai_pack(
-        compiled_models=compiled_models,
-        ai_pack_dir=ai_pack_dir,
-        ai_pack_name=ai_pack_name + '_mtk',
-        litert_model_name=litert_model_name + '_mtk',
-    )
-
-
-def _export_model_files_to_mtk_ai_pack(
-    compiled_models: types.CompilationResult,
-    ai_pack_dir: pathlib.Path,
-    ai_pack_name: str,
-    litert_model_name: str,
-):
-  """Exports the model tflite files to the MTK AI pack directory structure."""
-  model_export_dir = ai_pack_dir / ai_pack_name / 'src/main/assets'
-  os.makedirs(model_export_dir, exist_ok=True)
-  for backend, model in compiled_models.models_with_backend:
-    if not _is_mobile_device_backend(backend):
-      continue
-    backend_id = backend.id()
-    target_id = backend.target_id
-    if backend_id == fallback_backend.FallbackBackend.id():
-      target_id = 'other'
-    elif backend_id == mediatek_backend.MediaTekBackend.id():
-      target_id = backend.target_id.replace(
-          mtk_target.SocManufacturer.MEDIATEK, 'Mediatek'
-      )
-    group_name = 'model#group_' + target_id
-    export_dir = model_export_dir / group_name
-    os.makedirs(export_dir, exist_ok=True)
-    if backend_id != mediatek_backend.MediaTekBackend.id():
-      # Skip non-MTK targets, just create a placeholder file.
-      placeholder_file = export_dir / 'placeholder.txt'
-      placeholder_file.touch()
-      continue
-    model_export_path = export_dir / (litert_model_name + common.DOT_TFLITE)
-    if not model.in_memory:
-      model.load()
-    model.save(model_export_path, export_only=True)
-
-
-def _build_targeting_config(compiled_backends: list[types.Backend]) -> str:
-  """Builds device-targeting-config in device_targeting_configuration.xml."""
-  device_groups = []
-  for backend in compiled_backends:
-    if not _is_mobile_device_backend(backend):
-      continue
-    target = backend.target
-    device_group = _target_to_ai_pack_info(target)
-    if device_group:
-      device_groups.append(device_group)
-  device_groups = '\n'.join(device_groups)
-  return _DEVICE_TARGETING_CONFIGURATION.format(device_groups=device_groups)
-
-
-def _target_to_ai_pack_info(target: types.Target) -> str | None:
-  """Builds the device group used in device_targeting_configuration.xml."""
-  if isinstance(target, qnn_target.Target):
-    group_name = str(target)
-    selector = _process_qnn_target(target)
-    device_selectors = [
-        _DEVICE_SELECTOR_TEMPLATE.format(soc_man=man, soc_model=model)
-        for man, model in selector
-    ]
-    device_selectors = '\n'.join(device_selectors)
-    device_group = _DEVICE_GROUP_TEMPLATE.format(
-        device_group_name=group_name, device_selectors=device_selectors
-    )
-    return device_group
-  elif isinstance(target, mtk_target.Target):
-    group_name = str(target).replace(
-        mtk_target.SocManufacturer.MEDIATEK, 'Mediatek'
-    )
-    # TODO: b/407453529 - Support MTK SDK Version / OS version in selector.
-    selector = _process_mtk_target(target)
-    device_selector = _DEVICE_SELECTOR_TEMPLATE.format(
-        soc_man=selector[0], soc_model=selector[1]
-    )
-    device_group = _DEVICE_GROUP_TEMPLATE.format(
-        device_group_name=group_name, device_selectors=device_selector
-    )
-    return device_group
-  elif isinstance(target, fallback_backend.FallbackTarget):
-    # Don't need to have device selector for fallback target.
-    return None
-  else:
-    print('unsupported target ', target)
-    return None
-
-
-# TODO: b/407453529 - Auto-generate this function from CSVs.
-def _process_qnn_target(target: qnn_target.Target) -> list[tuple[str, str]]:
-  """Returns the list of (manufacturer, model) for the given QNN target."""
-  # Play cannot distinguish between Qualcomm and QTI for now.
-  manufacturer = ['Qualcomm', 'QTI']
-  models = [str(target.soc_model)]
-  return list(itertools.product(manufacturer, models))
-
-
-# TODO: b/407453529 - Auto-generate this function from CSVs.
-def _process_mtk_target(
-    target: mtk_target.Target,
-) -> tuple[str, str]:
-  """Returns tuple of (manufacturer, model) for the given MTK target."""
-  # Play cannot distinguish between Qualcomm and QTI for now.
-  return str(target.soc_manufacturer).replace(
-      mtk_target.SocManufacturer.MEDIATEK, 'Mediatek'
-  ), str(target.soc_model)
-
-
-def _write_targeting_config(
-    compiled_models: types.CompilationResult, ai_pack_dir: pathlib.Path
-) -> None:
-  """Writes device_targeting_configuration.xml for the given compiled models."""
-  compiled_backends = [x for x, _ in compiled_models.models_with_backend]
-  targeting_config = _build_targeting_config(
-      compiled_backends=compiled_backends
-  )
-
-  targeting_config_path = ai_pack_dir / 'device_targeting_configuration.xml'
-  targeting_config_path.write_text(targeting_config)
-
-
-def export(
-    compiled_models: types.CompilationResult,
-    ai_pack_dir: pathlib.Path | str,
-    ai_pack_name: str,
-    litert_model_name: str,
-) -> None:
-  """Exports the compiled models to AI pack format.
-
-  This function will export the compiled models to corresponding directory
-  structure:
-
-  {ai_pack_dir}/
-    AiPackManifest.xml
-    device_targeting_configuration.xml
-    {ai_pack_name}/src/main/assets/
-      model#group_target_1/
-        {litert_model_name}.tflite
-      model#group_target_2/
-        {litert_model_name}.tflite
-      model#group_target_3/
-        {litert_model_name}.tflite
-      model#group_other/
-        {litert_model_name}.tflite
-
-  Args:
-    compiled_models: The compiled models to export.
-    ai_pack_dir: The directory to export the AI pack to.
-    ai_pack_name: The name of the AI pack.
-    litert_model_name: The name of the model in the litert format.
-  """
-  if isinstance(ai_pack_dir, str):
-    ai_pack_dir = pathlib.Path(ai_pack_dir)
-
-  ai_pack_dir.mkdir(parents=True, exist_ok=True)
-
-  _export_model_files_to_ai_pack(
-      compiled_models=compiled_models,
-      ai_pack_dir=ai_pack_dir,
-      ai_pack_name=ai_pack_name,
-      litert_model_name=litert_model_name,
-  )
-  _write_targeting_config(
-      compiled_models=compiled_models, ai_pack_dir=ai_pack_dir
-  )

ai_edge_litert/aot/aot_compile.py

@@ -1,152 +0,0 @@
-# Copyright 2025 The LiteRT Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""AOT Compilation for LiteRT model."""
-import pathlib
-import tempfile
-
-from ai_edge_litert.aot import prepare_for_npu as core
-from ai_edge_litert.aot.core import apply_plugin
-from ai_edge_litert.aot.core import components
-from ai_edge_litert.aot.core import mlir_transforms
-from ai_edge_litert.aot.core import types
-from ai_edge_litert.aot.vendors import import_vendor
-
-
-def aot_compile(
-    input_model: types.Model | str,
-    output_dir: str | pathlib.Path | None = None,
-    target: types.Target | list[types.Target] | None = None,
-    config: (
-        types.CompilationConfig | list[types.CompilationConfig] | None
-    ) = None,
-    quantizer: components.AieQuantizerT | None = None,
-    keep_going: bool = True,
-    subgraphs_to_compile: list[int] | None = None,
-    **kwargs,
-) -> types.CompilationResult:
-  """Prepares a TFLite model for NPU execution.
-
-  High level command that erforms various backend specific pre-processing steps
-  and then applies an NPU compiler to the given model.
-
-  Args:
-    input_model: The input model to compile.
-    output_dir: Directory to write the output files to. If not specified, the
-      output files will be written to the same directory as the input file.
-    target: The target to compile for. If not specified, will compile to all
-      registered targets.
-    config: The compilation config(s). Cannot be specified with target.
-    quantizer: The quantizer to use for quantization.
-    keep_going: Whether to keep going if some backends fail. If False, fail
-      fast on the first error and raise an exception.
-    subgraphs_to_compile: The subgraph index list to compile to NPU. If None,
-      compile all subgraphs.
-    **kwargs: Additional arguments to pass to the backend.
-
-  Returns:
-    Compiled models.
-  """
-  # Only one of target or config is needed.
-  if target and config:
-    raise ValueError("Cannot specify both target and config.")
-
-  if config is None:
-    if target is None:
-      target = import_vendor.AllRegisteredTarget()
-    if isinstance(target, types.Target):
-      config = types.CompilationConfig(target=target)
-    elif isinstance(target, list):
-      config = [types.CompilationConfig(target=t) for t in target]
-    else:
-      raise ValueError("Unsupported target type.")
-
-  if isinstance(input_model, str):
-    input_path = pathlib.Path(input_model)
-    input_model = types.Model.create_from_path(input_path)
-
-  # Resolve output paths.
-  temp_dir = None
-  if not output_dir:
-    if input_model.in_memory:
-      # Use a temp dir for in-memory models.
-      # The temp dir will be cleaned up after the models are compiled and loaded
-      # back to memory (i.e. function returns).
-      temp_dir = tempfile.TemporaryDirectory()
-      output_dir = temp_dir.name
-    else:
-      input_path = input_model.path
-      output_dir = input_path.parent / "_compiled_models"
-      output_dir.mkdir(parents=True, exist_ok=True)
-      output_dir = str(output_dir)
-  output_dir_path = pathlib.Path(output_dir)
-
-  if isinstance(config, types.CompilationConfig) or not config:
-    if config:
-      # Make pytype happy.
-      assert isinstance(config, types.CompilationConfig)
-      kw_config = config.to_dict() | kwargs
-    else:
-      kw_config = kwargs
-
-    backend_class = core.resolve_backend(kw_config)
-
-    quant_recipe = kw_config.get("quantize_recipe", None)
-    if quant_recipe:
-      assert quantizer is not None, "Quantizer is required for quantization."
-
-    results = core.prepare_for_npu(
-        input_model,
-        output_dir_path,
-        backend_class,
-        kw_config,
-        transforms=mlir_transforms.MlirTransforms(),
-        quantizer=quantizer,
-        plugin=apply_plugin.ApplyPlugin(
-            experimental_capture_stderr=True,
-            subgraphs_to_compile=subgraphs_to_compile,
-        ),
-        keep_going=keep_going,
-    )
-  elif isinstance(config, list):
-    kw_configs = [c.to_dict() | kwargs for c in config]
-
-    configs_with_backend = [(core.resolve_backend(c), c) for c in kw_configs]
-    requires_quantizer = any("quantize_recipe" in c for c in kw_configs)
-    if requires_quantizer and quantizer is None:
-      raise ValueError("Quantizer is required for quantization.")
-
-    results = core.prepare_for_npu_multiple_configs(
-        input_model,
-        output_dir_path,
-        configs_with_backend,
-        transforms=mlir_transforms.MlirTransforms(),
-        quantizer=quantizer,
-        plugin=apply_plugin.ApplyPlugin(
-            experimental_capture_stderr=True,
-            subgraphs_to_compile=subgraphs_to_compile,
-        ),
-        keep_going=keep_going,
-    )
-  else:
-    # Should not reach here.
-    raise ValueError("Unsupported config type.")
-
-  if temp_dir:
-    # Load the models to memory before cleaning up the temp dir.
-    results.load()
-    temp_dir.cleanup()
-
-  return results

ai_edge_litert/aot/core/apply_plugin.py

@@ -1,146 +0,0 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""Wrapper for calling the apply plugin tooling."""
-
-
-import os
-import pathlib
-import re
-import subprocess
-import tempfile
-
-from ai_edge_litert.aot.core import common
-from ai_edge_litert.aot.core import components
-from ai_edge_litert.aot.core import types
-
-_BINARY = pathlib.Path("tools/apply_plugin_main")
-
-_RE_PARTITION_STATS = re.compile(
-    r"PartitionSubgraph: (\d+), selected num ops: (\d+), from totoal ops:"
-    r" (\d+), num partitions: (\d+)"
-)
-
-
-class ApplyPlugin(components.ApplyPluginT):
-  """Wrapper for calling the apply plugin tooling."""
-
-  def __init__(
-      self,
-      experimental_capture_stderr: bool = False,
-      subgraphs_to_compile: list[int] | None = None,
-  ):
-    self._experimental_capture_stderr = experimental_capture_stderr
-    self._subgraphs_to_compile = subgraphs_to_compile
-
-  @property
-  def default_err(self) -> str:
-    # NOTE: Capture stderr from underlying binary.
-    return "--"
-
-  @property
-  def component_name(self) -> str:
-    return "apply_plugin"
-
-  def __call__(
-      self,
-      input_model: types.Model,
-      output_model: types.Model,
-      soc_manufacturer: str,
-      soc_model: str,
-      sdk_libs_path: str | None = None,
-      **kwargs,
-  ):
-    """Applies a plugin to the input model.
-
-    Args:
-      input_model: The path to the input model.
-      output_model: The path to the output model.
-      soc_manufacturer: The SOC manufacturer of the plugin.
-      soc_model: The SOC model of the plugin.
-      sdk_libs_path: The path to the SDK libs. If not provided,
-        the default SDK path will be used.
-      **kwargs: Additional arguments to pass to the underlying binary.
-
-    Returns:
-      The output model.
-
-    Raises:
-      ValueError: If no tflite model was created by the underying binary.
-    """
-    if input_model.in_memory:
-      tmp_file = tempfile.NamedTemporaryFile(mode="wb")
-      input_model.save(tmp_file.name)
-    else:
-      tmp_file = None
-
-    binary = common.get_resource(_BINARY)
-    args = [
-        str(binary),
-        "--cmd=apply",
-        f"--model={str(input_model.path)}",
-        f"--o={str(output_model.path)}",
-        f"--soc_manufacturer={soc_manufacturer}",
-        f"--soc_model={soc_model}",
-        f"--err={self.default_err}",
-    ]
-    extra_args = [f"--{key}={value}" for key, value in kwargs.items()]
-    args.extend(extra_args)
-    if self._subgraphs_to_compile:
-      subgraphs_to_compile = ",".join(
-          str(s) for s in self._subgraphs_to_compile
-      )
-      args.append(f"--subgraphs={subgraphs_to_compile}")
-    env = os.environ.copy()
-    ld_library_path = common.construct_ld_library_path()
-    if ld_library_path:
-      if sdk_libs_path:
-        ld_library_path = f"{sdk_libs_path}{os.pathsep}{ld_library_path}"
-      env["LD_LIBRARY_PATH"] = ld_library_path
-
-    result = subprocess.run(
-        args,
-        check=False,
-        text=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        env=env,
-    )
-    if result.returncode:
-      log_file = tempfile.NamedTemporaryFile(mode="w", delete=False)
-      log_file.write(result.stdout)
-      log_file.close()
-      raise ValueError(
-          f"{self.component_name} failed to apply plugin. See"
-          f" {log_file.name} for details."
-      )
-
-    if not common.is_tflite(output_model.path):
-      raise ValueError(f"{output_model.path} is not a TFLite model.")
-
-    partition_stats = _RE_PARTITION_STATS.findall(result.stdout)
-    output_model.partition_stats = types.PartitionStats(
-        subgraph_stats=[
-            types.SubgraphPartitionStats(
-                subgraph_index=int(s[0]),
-                num_ops_offloaded=int(s[1]),
-                num_total_ops=int(s[2]),
-                num_partitions_offloaded=int(s[3]),
-            )
-            for s in partition_stats
-        ]
-    )
-    if tmp_file is not None:
-      tmp_file.close()