litert-cli 0.1.0__py3-none-any.whl

litert_cli/commands/quantize.py

@@ -0,0 +1,193 @@
+# Copyright 2026 The LiteRT CLI Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Quantize a LiteRT model using AI Edge Quantizer."""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import textwrap
+
+import click
+from litert_cli.core import deps
+
+
+@click.command(
+    "quantize",
+    help=textwrap.dedent("""\
+        Quantize a LiteRT model.
+
+        MODEL: Path to the input .tflite model.
+
+        Examples:
+
+          Dynamic INT8 Quantization (Default):
+
+            $ litert quantize raw_model.tflite --output quant_model.tflite
+
+          Weight-Only INT8 Quantization:
+
+            $ litert quantize raw_model.tflite --output quant_model.tflite \
+                --recipe weight_only_wi8_afp32
+
+          Static Quantization (Requires calibration data):
+
+            $ litert quantize raw_model.tflite --recipe static_wi8_ai8 \
+                --calibration-data calib_data.py --output quant_model.tflite
+
+          Custom Recipe:
+
+            $ litert quantize raw_model.tflite --custom-recipe recipe.json \
+                --output quant_model.tflite
+        """),
+)
+@click.argument("model", type=str)
+@click.option(
+    "--output",
+    type=click.Path(dir_okay=False, writable=True, path_type=pathlib.Path),
+    required=False,
+    default=None,
+    help="Path to save the output quantized .tflite model.",
+)
+@click.option(
+    "--recipe",
+    "quant_recipe",
+    type=str,
+    default="dynamic_wi8_afp32",
+    help=(
+        "Built-in quantization recipe to apply (e.g., dynamic_wi8_afp32,"
+        " weight_only_wi8_afp32, static_wi8_ai8). See"
+        " 'ai_edge_quantizer.recipe' for full list."
+    ),
+)
+@click.option(
+    "--calibration-data",
+    type=click.Path(
+        exists=True, dir_okay=False, resolve_path=True, path_type=pathlib.Path
+    ),
+    help=(
+        "Path to Python script providing calibration data (required for"
+        " 'static_wi8_ai8')."
+    ),
+)
+@click.option(
+    "--custom-recipe",
+    type=click.Path(
+        exists=True, dir_okay=False, resolve_path=True, path_type=pathlib.Path
+    ),
+    help="Path to JSON recipe file for custom configurations.",
+)
+@deps.require_extra("quantize")
+def quantize_cmd(
+    model: str,
+    output: pathlib.Path | None,
+    quant_recipe: str,
+    calibration_data: pathlib.Path | None,
+    custom_recipe: pathlib.Path | None,
+) -> None:
+  r"""Quantize a LiteRT model.
+
+  Args:
+    model: Path to the input .tflite model or Model Reference.
+    output: Path to save the output quantized .tflite model.
+    quant_recipe: Built-in quantization recipe to apply.
+    calibration_data: Path to Python script providing calibration data.
+    custom_recipe: Path to JSON recipe file for custom configurations.
+
+  Raises:
+    click.UsageError: If calibration data is missing when required.
+    click.ClickException: If calibration script loading fails or is invalid.
+  """
+  from ai_edge_quantizer import quantizer as aeq  # pylint: disable=g-import-not-at-top
+  from ai_edge_quantizer import qtyping  # pylint: disable=g-import-not-at-top
+
+  from litert_cli.core import models as core_models
+
+  resolved_model_path, _ = core_models.resolve_model_reference(model)
+
+  if str(resolved_model_path) != str(model):
+    click.echo(f"Resolved model '{model}' to '{resolved_model_path}'")
+
+  model_path = pathlib.Path(resolved_model_path)
+
+  if quant_recipe == "static_wi8_ai8" and calibration_data is None:
+    raise click.UsageError(
+        "--calibration-data is required when --recipe is 'static_wi8_ai8'."
+    )
+
+  resolved_output = output or model_path.with_name(
+      f"{model_path.stem}_quant.tflite"
+  )
+
+  click.echo(f"Quantizing '{model_path}' to '{resolved_output}'...")
+  quantizer = aeq.Quantizer(str(model_path))
+
+  if custom_recipe:
+    click.echo(f"Loading custom recipe from '{custom_recipe}'...")
+    import json
+
+    with open(custom_recipe, "r") as f:
+      recipe_content = json.load(f)
+    quantizer.load_quantization_recipe(recipe_content)
+  else:
+    from ai_edge_quantizer import recipe as aeq_recipe
+
+    if hasattr(aeq_recipe, quant_recipe):
+      click.echo(f"Loading built-in recipe '{quant_recipe}'...")
+      recipe_obj = getattr(aeq_recipe, quant_recipe)()
+      quantizer.load_quantization_recipe(recipe_obj)
+    else:
+      click.echo(
+          f"Fallback configuring dynamic quantization for '{quant_recipe}'..."
+      )
+      quantizer.add_dynamic_config(
+          regex=".*",
+          operation_name=qtyping.TFLOperationName.ALL_SUPPORTED,
+          num_bits=8,
+      )
+
+  calibration_result = None
+  if quantizer.need_calibration:
+    if calibration_data is None:
+      raise click.UsageError(
+          "Calibration data is required for the specified recipe. Use "
+          "--calibration-data <script.py>"
+      )
+    click.echo(f"Loading calibration data from '{calibration_data}'...")
+
+    spec = importlib.util.spec_from_file_location(
+        "calib_module", str(calibration_data)
+    )
+    if spec is None or spec.loader is None:
+      raise click.ClickException(
+          f"Failed to load calibration script from '{calibration_data}'"
+      )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    if not hasattr(module, "get_calibration_data"):
+      raise click.ClickException(
+          "Calibration script must define 'get_calibration_data()' function."
+      )
+    calib_data = module.get_calibration_data()
+    click.echo("Running calibration...")
+    calibration_result = quantizer.calibrate(calib_data)
+
+  click.echo("Applying quantization...")
+  result = quantizer.quantize(calibration_result)
+  resolved_output.parent.mkdir(parents=True, exist_ok=True)
+  result.export_model(str(resolved_output), overwrite=True)
+
+  click.secho(f"Quantization complete! Saved to {resolved_output}", fg="green")

litert_cli/commands/run/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2026 The LiteRT CLI Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Empty init for run subpackage."""

litert_cli/commands/run/android.py

@@ -0,0 +1,394 @@
+# Copyright 2026 The LiteRT CLI Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Android execution engine for LiteRT models.
+
+Use adb to push the model and the run_model binary to the device, then run the
+model on the device.
+
+Usage Examples:
+  1. Run a model on an Android device:
+     $ litert run /path/to/model.tflite --android
+
+  2. Run with NPU acceleration and CPU fallback:
+     $ litert run /path/to/model.tflite --android --npu --cpu
+     OR
+     $ litert run /path/to/model.tflite --android --accelerator npu,cpu
+
+  3. Run with custom inputs:
+     $ litert run /path/to/model.tflite --android --input input_name=value
+
+  4. Run with multiple inputs:
+     $ litert run /path/to/model.tflite --android --input input1=value1 --input
+     input2=value2
+
+  5. Run with specific signature:
+     $ litert run /path/to/model.tflite --android --signature_index 0
+
+  6. Run with multiple iterations:
+     $ litert run /path/to/model.tflite --android --iterations 10
+
+  7. Print tensor details:
+     $ litert run /path/to/model.tflite --android --print-tensors
+
+  8. Run with sample size:
+     $ litert run /path/to/model.tflite --android --sample-size 100
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import pathlib
+import shlex
+import subprocess
+import tempfile
+import time
+
+import click
+from litert_cli.core import android_utils
+from litert_cli.core import constants
+from litert_cli.core import inputs as inputs_utils
+from litert_cli.core import npu_utils as npu
+
+
+def _prepare_inputs_on_device(
+    *,
+    model_path: pathlib.Path,
+    inputs: Sequence[str],
+    signature_index: int,
+    android_root: str,
+) -> str:
+  """Parses inputs locally and pushes files to device, returns remote dir.
+
+  Args:
+    model_path: Local path to the LiteRT model file.
+    inputs: Tuple of input assignments (e.g. 'name=value').
+    signature_index: Signature index to run on device.
+    android_root: Remote workspace root path on Android device.
+
+  Returns:
+    Remote directory path where inputs are uploaded, or empty if no inputs.
+
+  Raises:
+    click.ClickException: If parsing fails or device push fails.
+  """
+  if not inputs:
+    return ""
+
+  try:
+    click.echo("Parsing inputs locally before pushing to device...")
+    from ai_edge_litert.compiled_model import CompiledModel  # pylint: disable=g-import-not-at-top
+
+    cm = CompiledModel.from_file(str(model_path))
+    signatures = cm.get_signature_list()
+    if not signatures:
+      click.secho("No signatures found in the model for inputs.", fg="yellow")
+      return ""
+
+    sig_info = cm.get_signature_by_index(signature_index)
+    sig_key = sig_info["key"]
+    input_details = cm.get_input_tensor_details(sig_key)
+
+    # 2. Process input strings mapping (e.g., name=value or literal value).
+    parsed_inputs = {}
+    for inp in inputs:
+      if "=" in inp:
+        k, v = inp.split("=", 1)
+        parsed_inputs[k] = v
+      else:
+        parsed_inputs["_default_"] = inp
+
+    has_inputs = False
+    # 3. Convert literals or files to raw binaries in a local temporary
+    # directory.
+    with tempfile.TemporaryDirectory() as temp_dir:
+      temp_path = pathlib.Path(temp_dir)
+      for name, details in input_details.items():
+        input_data_str = parsed_inputs.get(name) or parsed_inputs.get(
+            "_default_"
+        )
+        if input_data_str:
+          shape = details.get("shape", [1])
+          tensor_type = details.get("dtype", "?")
+          np_dtype = inputs_utils.get_np_dtype(tensor_type)
+
+          click.echo(
+              f"  Preparing input {name!r} from {input_data_str!r} (shape:"
+              f" {shape}, dtype: {tensor_type})"
+          )
+          data = inputs_utils.parse_input(input_data_str, shape, np_dtype)
+
+          raw_file_path = temp_path / f"{name}.raw"
+          data.tofile(raw_file_path)
+          has_inputs = True
+
+      if has_inputs:
+        # 4. Create remote directory and push input data batch tree to device.
+        remote_input_dir = f"{android_root}/inputs_{int(time.time())}"
+        click.echo(f"Pushing processed inputs to {remote_input_dir}...")
+        subprocess.run(
+            ["adb", "shell", f"mkdir -p {shlex.quote(remote_input_dir)}"],
+            check=True,
+        )
+        for local_file in temp_path.iterdir():
+          subprocess.run(
+              [
+                  "adb",
+                  "push",
+                  local_file,
+                  f"{remote_input_dir}/{local_file.name}",
+              ],
+              check=True,
+              stdout=subprocess.DEVNULL,
+              stderr=subprocess.DEVNULL,
+          )
+        return remote_input_dir
+
+    return ""
+
+  except Exception as e:  # pylint: disable=broad-exception-caught
+    raise click.ClickException(
+        f"Failed to prepare inputs for Android model {model_path}: {e!r}"
+    ) from e
+
+
+def run_android(
+    *,
+    model_path: str,
+    inputs: Sequence[str],
+    accelerator: str,
+    signature_index: int,
+    iterations: int,
+    print_tensors: bool,
+    sample_size: int,
+) -> None:
+  """Runs the model on an attached Android device using adb and run_model.
+
+  Args:
+    model_path: Local path to the LiteRT model file (.tflite).
+    inputs: Tuple of input assignments (e.g. 'name=value').
+    accelerator: Hardware accelerator ('cpu', 'gpu', 'npu').
+    signature_index: Signature index to execute.
+    iterations: Number of execute loops for remote runner.
+    print_tensors: Whether to print remote stats after execution completes.
+    sample_size: Limit execution sample stream print length per tensor.
+
+  Raises:
+    click.ClickException: On device error setup or failed execution triggers.
+  """
+  accel_list = [a.strip().lower() for a in accelerator.split(",") if a.strip()]
+  click.echo("Preparing to run on Android device via adb...")
+  android_utils.check_adb()
+
+  # Set up Android working directory
+  android_root = constants.LITERT_CLI_ANDROID_ROOT
+  try:
+    subprocess.run(
+        ["adb", "shell", f"mkdir -p {shlex.quote(android_root)}"],
+        check=True,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+  except subprocess.CalledProcessError as e:
+    raise click.ClickException(
+        f"Failed to create directory {android_root} on device: {e!r}"
+    ) from e
+
+  # Create remote execution tracking paths
+  model_name = pathlib.Path(model_path).name
+  remote_model_path = f"{android_root}/{model_name}"
+
+  # Determine device ABI
+  abi = android_utils.get_android_abi()
+  click.echo(f"Detected Android device ABI: {abi}")
+
+  run_model_bin = android_utils.find_android_binary("run_model", abi)
+
+  # Download libraries
+  lib_litert = android_utils.find_android_lib("libLiteRt.so", abi)
+  lib_clgl = android_utils.find_android_lib("libLiteRtClGlAccelerator.so", abi)
+
+  # Push model file and runner binary to Android device
+  click.echo(f"Pushing model {model_name} to device...")
+  subprocess.run(["adb", "push", model_path, remote_model_path], check=True)
+
+  remote_run_model_path = f"{android_root}/run_model"
+  if (
+      subprocess.run(
+          ["adb", "shell", f"[ -f {shlex.quote(remote_run_model_path)} ]"],
+          check=False,
+      ).returncode
+      == 0
+  ):
+    click.echo("  Skipping run_model (already on device)")
+  else:
+    click.echo("Pushing run_model to device...")
+    subprocess.run(
+        ["adb", "push", run_model_bin, remote_run_model_path], check=True
+    )
+
+  # Push libraries to default path
+  remote_lib_litert = f"{android_root}/{lib_litert.name}"
+  if (
+      subprocess.run(
+          ["adb", "shell", f"[ -f {shlex.quote(remote_lib_litert)} ]"],
+          check=False,
+      ).returncode
+      == 0
+  ):
+    click.echo(f"  Skipping {lib_litert.name} (already on device)")
+  else:
+    click.echo(f"Pushing {lib_litert.name} to device...")
+    subprocess.run(["adb", "push", lib_litert, remote_lib_litert], check=True)
+
+  remote_lib_clgl = f"{android_root}/{lib_clgl.name}"
+  if (
+      subprocess.run(
+          ["adb", "shell", f"[ -f {shlex.quote(remote_lib_clgl)} ]"],
+          check=False,
+      ).returncode
+      == 0
+  ):
+    click.echo(f"  Skipping {lib_clgl.name} (already on device)")
+  else:
+    click.echo(f"Pushing {lib_clgl.name} to device...")
+    subprocess.run(["adb", "push", lib_clgl, remote_lib_clgl], check=True)
+
+  remote_input_dir = _prepare_inputs_on_device(
+      model_path=pathlib.Path(model_path),
+      inputs=inputs,
+      signature_index=signature_index,
+      android_root=android_root,
+  )
+
+  # Pass None as device_id to use the default connected device.
+  remote_dispatch_dir = (
+      npu.push_npu_runtime_libraries(None, android_root)
+      if "npu" in accel_list
+      else ""
+  )
+
+  if "npu" in accel_list:
+    # Download and push SOC-specific LiteRT dispatch and compiler plugin libraries
+    target_model = npu.get_soc_target_model(None)
+    soc_vendor = "mediatek" if "mt" in target_model else "qualcomm"
+    lib_dispatch = android_utils.find_npu_dispatch_lib(soc_vendor, abi)
+    lib_compiler = android_utils.find_npu_compiler_plugin_lib(soc_vendor, abi)
+
+    remote_lib_dispatch = f"{android_root}/{lib_dispatch.name}"
+    if (
+        subprocess.run(
+            ["adb", "shell", f"[ -f {shlex.quote(remote_lib_dispatch)} ]"],
+            check=False,
+        ).returncode
+        == 0
+    ):
+      click.echo(f"  Skipping {lib_dispatch.name} (already on device)")
+    else:
+      click.echo(f"Pushing {lib_dispatch.name} to device...")
+      subprocess.run(
+          ["adb", "push", lib_dispatch, remote_lib_dispatch], check=True
+      )
+
+    remote_lib_compiler = f"{android_root}/{lib_compiler.name}"
+    if (
+        subprocess.run(
+            ["adb", "shell", f"[ -f {shlex.quote(remote_lib_compiler)} ]"],
+            check=False,
+        ).returncode
+        == 0
+    ):
+      click.echo(f"  Skipping {lib_compiler.name} (already on device)")
+    else:
+      click.echo(f"Pushing {lib_compiler.name} to device...")
+      subprocess.run(
+          ["adb", "push", lib_compiler, remote_lib_compiler], check=True
+      )
+
+  click.echo("Executing on device...\n")
+
+  # Need to make the binary executable before running the model
+  subprocess.run(
+      ["adb", "shell", f"chmod +x {shlex.quote(remote_run_model_path)}"],
+      check=True,
+  )
+
+  # Forward valid flags
+  run_cmd_args = [remote_run_model_path, f"--graph={remote_model_path}"]
+  if accelerator != "cpu":
+    run_cmd_args.append(f"--accelerator={accelerator}")
+  if remote_dispatch_dir:
+    run_cmd_args.append(f"--dispatch_library_dir={remote_dispatch_dir}")
+    run_cmd_args.append(f"--compiler_plugin_library_dir={remote_dispatch_dir}")
+  if iterations > 1:
+    run_cmd_args.append(f"--iterations={iterations}")
+  if signature_index != 0:
+    run_cmd_args.append(f"--signature_index={signature_index}")
+  if print_tensors:
+    run_cmd_args.append("--print_tensors=true")
+  run_cmd_args.append(f"--sample_size={sample_size}")
+  if remote_input_dir:
+    run_cmd_args.append(f"--input_dir={remote_input_dir}")
+
+  try:
+    if remote_dispatch_dir:
+      env_vars = (
+          f"LD_LIBRARY_PATH={shlex.quote(f'{remote_dispatch_dir}:{android_root}')}"
+          f" ADSP_LIBRARY_PATH={shlex.quote(remote_dispatch_dir)}"
+      )
+    else:
+      env_vars = ""
+    cmd_str = f"{env_vars} " if env_vars else ""
+    cmd_str += " ".join(shlex.quote(arg) for arg in run_cmd_args)
+    process = subprocess.Popen(
+        ["adb", "shell", cmd_str],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+    )
+
+    from litert_cli.core.log_filters import RunLogFilter
+
+    output_lines = []
+    log_filter = RunLogFilter(constants.DEFAULT_QUIET, print_tensors)
+
+    for line in process.stdout:
+      output_lines.append(line)
+      if log_filter.should_show(line):
+        click.echo(line, nl=False)
+
+    process.wait()
+    if process.returncode != 0:
+      click.secho(
+          f"Execution failed on device with exit code {process.returncode}",
+          fg="red",
+      )
+      click.echo("Full output for debugging:")
+      for line in output_lines:
+        click.echo(line, nl=False)
+      raise click.ClickException("Execution failed on device.")
+  except Exception as e:
+    raise click.ClickException(f"Failed to execute on device: {e}")
+  finally:
+    # Cleanup remote paths
+    click.echo("Clearing remote files...")
+    cleanup_cmds = [
+        f"rm -f {shlex.quote(remote_model_path)}"
+        f" {shlex.quote(remote_run_model_path)}"
+    ]
+    if remote_input_dir:
+      cleanup_cmds.append(f"rm -rf {shlex.quote(remote_input_dir)}")
+    cleanup_cmd = " && ".join(cleanup_cmds)
+    subprocess.run(["adb", "shell", cleanup_cmd], check=False)