ai-edge-litert-nightly 2.0.4.dev20251102__cp313-cp313-macosx_12_0_arm64.whl

ai_edge_litert/__init__.py

@@ -0,0 +1 @@
+__version__ = "2.0.4.dev20251102"

ai_edge_litert/any_pb2.py

@@ -0,0 +1,37 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: google/protobuf/any.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'google/protobuf/any.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19google/protobuf/any.proto\x12\x0fgoogle.protobuf\"&\n\x03\x41ny\x12\x10\n\x08type_url\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c\x42v\n\x13\x63om.google.protobufB\x08\x41nyProtoP\x01Z,google.golang.org/protobuf/types/known/anypb\xa2\x02\x03GPB\xaa\x02\x1eGoogle.Protobuf.WellKnownTypesb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.protobuf.any_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n\023com.google.protobufB\010AnyProtoP\001Z,google.golang.org/protobuf/types/known/anypb\242\002\003GPB\252\002\036Google.Protobuf.WellKnownTypes'
+  _globals['_ANY']._serialized_start=46
+  _globals['_ANY']._serialized_end=84
+# @@protoc_insertion_point(module_scope)

ai_edge_litert/api_pb2.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: google/protobuf/api.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'google/protobuf/api.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import source_context_pb2 as google_dot_protobuf_dot_source__context__pb2
+from google.protobuf import type_pb2 as google_dot_protobuf_dot_type__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19google/protobuf/api.proto\x12\x0fgoogle.protobuf\x1a$google/protobuf/source_context.proto\x1a\x1agoogle/protobuf/type.proto\"\x81\x02\n\x03\x41pi\x12\x0c\n\x04name\x18\x01 \x01(\t\x12(\n\x07methods\x18\x02 \x03(\x0b\x32\x17.google.protobuf.Method\x12(\n\x07options\x18\x03 \x03(\x0b\x32\x17.google.protobuf.Option\x12\x0f\n\x07version\x18\x04 \x01(\t\x12\x36\n\x0esource_context\x18\x05 \x01(\x0b\x32\x1e.google.protobuf.SourceContext\x12&\n\x06mixins\x18\x06 \x03(\x0b\x32\x16.google.protobuf.Mixin\x12\'\n\x06syntax\x18\x07 \x01(\x0e\x32\x17.google.protobuf.Syntax\"\xd5\x01\n\x06Method\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x18\n\x10request_type_url\x18\x02 \x01(\t\x12\x19\n\x11request_streaming\x18\x03 \x01(\x08\x12\x19\n\x11response_type_url\x18\x04 \x01(\t\x12\x1a\n\x12response_streaming\x18\x05 \x01(\x08\x12(\n\x07options\x18\x06 \x03(\x0b\x32\x17.google.protobuf.Option\x12\'\n\x06syntax\x18\x07 \x01(\x0e\x32\x17.google.protobuf.Syntax\"#\n\x05Mixin\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04root\x18\x02 \x01(\tBv\n\x13\x63om.google.protobufB\x08\x41piProtoP\x01Z,google.golang.org/protobuf/types/known/apipb\xa2\x02\x03GPB\xaa\x02\x1eGoogle.Protobuf.WellKnownTypesb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.protobuf.api_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'\n\023com.google.protobufB\010ApiProtoP\001Z,google.golang.org/protobuf/types/known/apipb\242\002\003GPB\252\002\036Google.Protobuf.WellKnownTypes'
+  _globals['_API']._serialized_start=113
+  _globals['_API']._serialized_end=370
+  _globals['_METHOD']._serialized_start=373
+  _globals['_METHOD']._serialized_end=586
+  _globals['_MIXIN']._serialized_start=588
+  _globals['_MIXIN']._serialized_end=623
+# @@protoc_insertion_point(module_scope)

ai_edge_litert/compiled_model.py

@@ -0,0 +1,250 @@
+# 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.
+
+"""Python wrapper for LiteRT compiled models."""
+
+import os
+from typing import Any, Dict, List
+
+# pylint: disable=g-import-not-at-top
+if not os.path.splitext(__file__)[0].endswith(
+    os.path.join("ai_edge_litert", "compiled_model")
+):
+  # This file is part of litert package.
+  from ai_edge_litert.litert_wrapper.compiled_model_wrapper import (
+      _pywrap_litert_compiled_model_wrapper as _cm,
+  )
+  from ai_edge_litert.litert_wrapper.compiled_model_wrapper.hardware_accelerator import HardwareAccelerator
+  from ai_edge_litert.litert_wrapper.tensor_buffer_wrapper.tensor_buffer import TensorBuffer
+else:
+  # This file is part of ai_edge_litert package.
+  from ai_edge_litert import _pywrap_litert_compiled_model_wrapper as _cm
+  from ai_edge_litert.hardware_accelerator import HardwareAccelerator
+  from ai_edge_litert.tensor_buffer import TensorBuffer
+# pylint: enable=g-import-not-at-top
+
+
+class CompiledModel:
+  """Python wrapper for the C++ CompiledModelWrapper.
+
+  This class provides methods to load, inspect, and execute machine learning
+  models using the LiteRT runtime.
+  """
+
+  def __init__(self, c_model_ptr):
+    """Initializes the CompiledModel with a C++ model pointer.
+
+    Args:
+      c_model_ptr: Pointer to the underlying C++ CompiledModelWrapper.
+    """
+    self._model = c_model_ptr  # Pointer to C++ CompiledModelWrapper
+
+  @classmethod
+  def from_file(
+      cls,
+      model_path: str,
+  ) -> "CompiledModel":
+    """Creates a CompiledModel from a model file.
+
+    Args:
+      model_path: Path to the model file.
+
+    Returns:
+      A new CompiledModel instance.
+    """
+    ptr = _cm.CreateCompiledModelFromFile(
+        model_path,
+        compiler_plugin_path="",
+        dispatch_library_path="",
+        hardware_accel=HardwareAccelerator.CPU,
+    )
+    return cls(ptr)
+
+  @classmethod
+  def from_buffer(
+      cls,
+      model_data: bytes,
+  ) -> "CompiledModel":
+    """Creates a CompiledModel from an in-memory buffer.
+
+    Args:
+      model_data: Model data as bytes.
+
+    Returns:
+      A new CompiledModel instance.
+    """
+    ptr = _cm.CreateCompiledModelFromBuffer(
+        model_data,
+        compiler_plugin_path="",
+        dispatch_library_path="",
+        hardware_accel=HardwareAccelerator.CPU,
+    )
+    return cls(ptr)
+
+  def get_signature_list(self) -> Dict[str, Dict[str, List[str]]]:
+    """Returns a dictionary of all available model signatures.
+
+    Returns:
+      Dictionary mapping signature names to their input/output specifications.
+    """
+    return self._model.GetSignatureList()
+
+  def get_signature_by_index(self, index: int) -> Dict[str, Any]:
+    """Returns signature information for the given index.
+
+    Args:
+      index: Index of the signature to retrieve.
+
+    Returns:
+      Dictionary containing signature information.
+    """
+    return self._model.GetSignatureByIndex(index)
+
+  def get_num_signatures(self) -> int:
+    """Returns the number of signatures in the model.
+
+    Returns:
+      Number of signatures.
+    """
+    return self._model.GetNumSignatures()
+
+  def get_signature_index(self, key: str) -> int:
+    """Returns the index for a signature name.
+
+    Args:
+      key: Name of the signature.
+
+    Returns:
+      Index of the signature, or -1 if not found.
+    """
+    return self._model.GetSignatureIndex(key)
+
+  def get_input_buffer_requirements(
+      self, input_index: int, signature_index: int = 0
+  ) -> Dict[str, Any]:
+    """Returns memory requirements for an input tensor.
+
+    Args:
+      input_index: Index of the input tensor.
+      signature_index: Index of the signature. Default is 0 (first signature).
+
+    Returns:
+      Dictionary with buffer requirements (size, alignment, etc.).
+    """
+    return self._model.GetInputBufferRequirements(signature_index, input_index)
+
+  def get_output_buffer_requirements(
+      self, output_index: int, signature_index: int = 0
+  ) -> Dict[str, Any]:
+    """Returns memory requirements for an output tensor.
+
+    Args:
+      output_index: Index of the output tensor.
+      signature_index: Index of the signature. Default is 0 (first signature).
+
+    Returns:
+      Dictionary with buffer requirements (size, alignment, etc.).
+    """
+    return self._model.GetOutputBufferRequirements(
+        signature_index, output_index
+    )
+
+  def create_input_buffer_by_name(
+      self, signature_key: str, input_name: str
+  ) -> TensorBuffer:
+    """Creates an input TensorBuffer for the specified signature and input name.
+
+    Args:
+      signature_key: Name of the signature.
+      input_name: Name of the input tensor.
+
+    Returns:
+      A TensorBuffer object for the specified input.
+    """
+    capsule = self._model.CreateInputBufferByName(signature_key, input_name)
+    return TensorBuffer(capsule)
+
+  def create_output_buffer_by_name(
+      self, signature_key: str, output_name: str
+  ) -> TensorBuffer:
+    """Creates an output TensorBuffer for the specified signature and output name.
+
+    Args:
+      signature_key: Name of the signature.
+      output_name: Name of the output tensor.
+
+    Returns:
+      A TensorBuffer object for the specified output.
+    """
+    capsule = self._model.CreateOutputBufferByName(signature_key, output_name)
+    return TensorBuffer(capsule)
+
+  def create_input_buffers(self, signature_index: int) -> List[TensorBuffer]:
+    """Creates TensorBuffers for all inputs of the specified signature.
+
+    Args:
+      signature_index: Index of the signature.
+
+    Returns:
+      List of TensorBuffer objects for all inputs.
+    """
+    capsule_list = self._model.CreateInputBuffers(signature_index)
+    return [TensorBuffer(c) for c in capsule_list]
+
+  def create_output_buffers(self, signature_index: int) -> List[TensorBuffer]:
+    """Creates TensorBuffers for all outputs of the specified signature.
+
+    Args:
+      signature_index: Index of the signature.
+
+    Returns:
+      List of TensorBuffer objects for all outputs.
+    """
+    capsule_list = self._model.CreateOutputBuffers(signature_index)
+    return [TensorBuffer(c) for c in capsule_list]
+
+  def run_by_name(
+      self,
+      signature_key: str,
+      input_map: Dict[str, TensorBuffer],
+      output_map: Dict[str, TensorBuffer],
+  ) -> None:
+    """Runs inference using the named signature and tensor maps.
+
+    Args:
+      signature_key: Name of the signature to execute.
+      input_map: Dictionary mapping input names to TensorBuffer objects.
+      output_map: Dictionary mapping output names to TensorBuffer objects.
+    """
+    # Convert TensorBuffer objects to raw capsules
+    capsule_input_map = {k: v.capsule for k, v in input_map.items()}
+    capsule_output_map = {k: v.capsule for k, v in output_map.items()}
+    self._model.RunByName(signature_key, capsule_input_map, capsule_output_map)
+
+  def run_by_index(
+      self,
+      signature_index: int,
+      input_buffers: List[TensorBuffer],
+      output_buffers: List[TensorBuffer],
+  ) -> None:
+    """Runs inference using the indexed signature and tensor lists.
+
+    Args:
+      signature_index: Index of the signature to execute.
+      input_buffers: List of input TensorBuffer objects.
+      output_buffers: List of output TensorBuffer objects.
+    """
+    input_capsules = [tb.capsule for tb in input_buffers]
+    output_capsules = [tb.capsule for tb in output_buffers]
+    self._model.RunByIndex(signature_index, input_capsules, output_capsules)