litert-cli 0.1.0__py3-none-any.whl

litert_cli/commands/visualize.py

@@ -0,0 +1,234 @@
+# 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.
+# ==============================================================================
+
+"""Command line interface for visualizing LiteRT models.
+
+This module provides the `litert visualize` command, which launches the
+Model Explorer web application to inspect the architecture of TFLite models.
+
+Key Features:
+- Background Execution: The web server runs entirely in the background,
+  instantly returning control of the terminal to the user.
+- URL Printing: The command calculates and prints a click-able URL directly
+  in the terminal, bypassing IDE port-forwarding interceptions that often
+  strip URL query parameters.
+- Server Reuse: By default, multiple visualizations reuse the same server
+  instance, preventing port exhaustion and reducing memory footprint.
+"""
+
+from __future__ import annotations
+
+import json
+import pathlib
+import socket
+import subprocess
+import sys
+import textwrap
+import urllib.parse
+
+import click
+
+from ..core import deps
+
+
+def _is_port_in_use(port_num: int) -> bool:
+  """Checks if a port is in use on localhost.
+
+  Args:
+    port_num: Port to check.
+
+  Returns:
+    True if the port is bound and in use, False otherwise.
+  """
+  with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+    return s.connect_ex(('localhost', port_num)) == 0
+
+
+def _find_available_port(start_port: int, max_attempts: int = 20) -> int:
+  """Finds the first available port on localhost starting from start_port.
+
+  Args:
+    start_port: Port to start searching from.
+    max_attempts: Maximum number of ports to check.
+
+  Returns:
+    The first available port found.
+  """
+  return next(
+      (
+          p
+          for p in range(start_port, start_port + max_attempts)
+          if not _is_port_in_use(p)
+      ),
+      start_port,
+  )
+
+
+@click.command(
+    'visualize',
+    help=textwrap.dedent("""\
+      Runs model explorer to visualize the model architecture.
+
+      MODEL_PATH: Path to the LiteRT model (.tflite) to visualize.
+
+      Examples:
+
+        1. Visualize a single model (starts a new server if none exists):
+
+          $ litert visualize /path/to/model.tflite
+
+        2. Visualize a different model (reuses existing server, refreshes browser):
+
+          $ litert visualize /path/to/another_model.tflite
+
+        3. Compare two models side-by-side (forces a new server on a new port):
+
+          $ litert visualize /path/to/model_A.tflite
+          $ litert visualize /path/to/model_B.tflite --no-reuse-server
+
+        4. Clean up and stop all running Model Explorer servers:
+
+          $ litert visualize --stop-all
+      """),
+)
+@click.argument(
+    'model_path',
+    type=str,
+    required=False,
+)
+@click.option(
+    '--reuse-server/--no-reuse-server',
+    default=True,
+    help='Reuse running Model Explorer server.',
+)
+@click.option(
+    '--stop-all',
+    is_flag=True,
+    help='Stop all running Model Explorer servers.',
+)
+@deps.require_extra('visualize')
+def visualize_cmd(
+    model_path: str | None, *, reuse_server: bool, stop_all: bool
+) -> None:
+  """Runs model explorer to visualize the model architecture.
+
+  Args:
+    model_path: Path to the LiteRT model (.tflite) or Model Reference to
+      visualize.
+    reuse_server: Whether to reuse an already running Model Explorer server.
+    stop_all: If True, stops all running Model Explorer servers instead.
+  """
+  from litert_cli.core import models as core_models
+
+  if stop_all:
+    click.echo('Attempting to stop all running Model Explorer servers...')
+    try:
+      # Use pkill to find and terminate all model-explorer processes
+      subprocess.run(
+          ['pkill', '-f', 'model-explorer'],
+          check=True,
+          stdout=subprocess.DEVNULL,
+          stderr=subprocess.DEVNULL,
+      )
+      click.secho(
+          'Successfully stopped all Model Explorer servers.',
+          fg='green',
+      )
+    except subprocess.CalledProcessError:
+      click.secho('No running Model Explorer servers found.', fg='yellow')
+    return
+
+  if not model_path:
+    raise click.UsageError('Missing argument "MODEL_PATH".')
+
+  resolved_model_path, _ = core_models.resolve_model_reference(model_path)
+  if str(resolved_model_path) != str(model_path):
+    click.echo(f"Resolved model '{model_path}' to '{resolved_model_path}'")
+
+  resolved_model_path = pathlib.Path(resolved_model_path).resolve()
+
+  click.echo(
+      f'Starting Model Explorer visualization for {resolved_model_path} in the'
+      ' background...'
+  )
+  # Check for available port so we can print the exact URL
+  port = (
+      8081
+      if reuse_server and _is_port_in_use(8081)
+      else _find_available_port(8081)
+  )
+
+  # Build the exact model explorer data URL
+  data = {'models': [{'url': str(resolved_model_path)}]}
+  data_param = urllib.parse.quote(
+      json.dumps(data, separators=(',', ':')), safe=':,'
+  )
+  url = f'http://localhost:{port}/?data={data_param}'
+
+  model_explorer_bin = None
+  if sys.executable:
+    python_dir = pathlib.Path(sys.executable).parent
+    # Assume model_explorer binary is in the same directory as the python
+    model_explorer_bin = python_dir / 'model-explorer'
+
+  cmd = [
+      str(model_explorer_bin)
+      if model_explorer_bin and model_explorer_bin.exists()
+      else 'model-explorer',
+      str(resolved_model_path),
+      '--no_open_in_browser',
+      f'--port={port}',
+  ]
+  if reuse_server:
+    cmd.append('--reuse_server')
+    cmd.append(f'--reuse_server_port={port}')
+
+  # Launch as a fully detached daemon so the terminal isn't blocked.
+  try:
+    subprocess.Popen(
+        cmd,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        start_new_session=True,
+    )
+
+    click.secho(
+        '\nModel Explorer server is running in the background.', fg='green'
+    )
+    click.secho(
+        '► Please click the link below to view your model:',
+        fg='cyan',
+        bold=True,
+    )
+    click.secho(f'\n    {url}\n')
+    if reuse_server:
+      click.secho(
+          'It will reuse the existing server instance if one is already'
+          ' running, avoiding port conflicts.',
+          fg='cyan',
+      )
+  except FileNotFoundError:
+    click.secho(
+        '\nError: "model-explorer" command not found.', fg='red', bold=True
+    )
+    click.secho(
+        'Please make sure Model Explorer is installed and available in your'
+        ' PATH.',
+        fg='yellow',
+    )
+    click.secho(
+        'You can install it via: pip install model-explorer',
+        fg='cyan',
+    )

litert_cli/core/android_utils.py

@@ -0,0 +1,304 @@
+# 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 utility functions for LiteRT CLI."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import pathlib
+import subprocess
+
+import click
+import requests
+
+from litert_cli.core import constants
+
+
+class DownloadError(Exception):
+  """A file download failed."""
+
+  pass
+
+
+def _construct_adb_command(command: Sequence[str], device_id: str | None) -> list[str]:
+  """Constructs an adb command list, including the device ID if provided."""
+  if device_id:
+    return ["adb", "-s", device_id, *command]
+  return ["adb", *command]
+
+
+def check_adb(device_id: str | None = None) -> None:
+  """Checks if adb device is available.
+
+  Verifies that exactly one authorized device is connected and responding to
+  commands.
+
+  Args:
+    device_id: Optional. The serial number of the target device.
+
+  Raises:
+    click.ClickException: If adb is missing or device is unauthorized/offline.
+  """
+  try:
+    subprocess.run(
+        _construct_adb_command(["get-state"], device_id),
+        check=True,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+  except FileNotFoundError as e:
+    raise click.ClickException(
+        "adb command not found. Please ensure Android platform-tools is in"
+        " your PATH."
+    ) from e
+  except subprocess.CalledProcessError as e:
+    device_info = f" for device {device_id}" if device_id else ""
+    raise click.ClickException(
+        f"No Android device found{device_info} or it is not authorized. Try"
+        " running 'adb devices' to check."
+    ) from e
+
+
+def get_android_abi(device_id: str | None = None) -> str:
+  """Gets the CPU ABI of the connected Android device via adb.
+
+  Args:
+    device_id: Optional. The serial number of the target device.
+
+  Returns:
+    The CPU ABI string (e.g., 'arm64-v8a').
+
+  Raises:
+    click.ClickException: If querying Android ABI fails or adb is missing.
+  """
+  try:
+    return subprocess.run(
+        _construct_adb_command(
+            ["shell", "getprop", "ro.product.cpu.abi"], device_id
+        ),
+        capture_output=True,
+        text=True,
+        check=True,
+    ).stdout.strip()
+  except subprocess.CalledProcessError as e:
+    device_info = f" for device {device_id}" if device_id else ""
+    raise click.ClickException(f"Error querying Android ABI{device_info}") from e
+
+
+def _ensure_downloaded_file(
+    abi: str,
+    file_name: str,
+    base_url: str,
+    *,
+    make_executable: bool = False,
+) -> pathlib.Path:
+  """Downloads a file for the given ABI if not cached.
+
+  Args:
+    abi: The Android CPU ABI (e.g., 'arm64-v8a').
+    file_name: The file name to download.
+    base_url: The base URL to download from.
+    make_executable: Whether to make the file executable.
+
+  Returns:
+    The absolute local path to the cached file.
+
+  Raises:
+    click.ClickException: If the architecture is not supported.
+    DownloadError: If the download or file operations fail.
+  """
+  if "arm64" not in abi:
+    raise click.ClickException(
+        f"Architecture {abi!r} is not supported for automatic downloading of"
+        f" {file_name!r}. Only arm64 is supported for these specific binaries."
+    )
+
+  download_url = f"{base_url}/{file_name}"
+  cache_dir = pathlib.Path(constants.LITERT_CLI_CACHE_DIR) / "binaries" / abi
+  cache_dir.mkdir(parents=True, exist_ok=True)
+
+  cached_file_path = cache_dir / file_name
+  if cached_file_path.exists():
+    return cached_file_path
+
+  click.secho(f"Downloading {file_name!r} for {abi!r}...", fg="cyan")
+  tmp_cached_file = cached_file_path.with_suffix(".tmp")
+  try:
+    with requests.get(download_url, stream=True, timeout=10) as response:
+      response.raise_for_status()
+      content_length = response.headers.get("Content-Length")
+      total_size = int(content_length) if content_length else 0
+
+      bar_length = total_size if total_size > 0 else None
+      bar_label = f"Downloading {file_name}"
+      if bar_length is None:
+        click.secho(
+            f"Content-Length header not found for {file_name!r}, using"
+            " indeterminate progress bar.",
+            fg="yellow",
+        )
+        bar_label += " (size unknown)"
+
+      with click.progressbar(
+          length=bar_length, label=bar_label
+      ) as bar:
+        with open(tmp_cached_file, "wb") as f:
+          for buffer in response.iter_content(chunk_size=8192):
+            f.write(buffer)
+            bar.update(len(buffer))
+
+    tmp_cached_file.rename(cached_file_path)
+    if make_executable:
+      cached_file_path.chmod(0o755)
+    return cached_file_path
+  except (requests.exceptions.RequestException, OSError) as e:
+    if cached_file_path.exists():
+      cached_file_path.unlink()
+    if tmp_cached_file.exists():
+      tmp_cached_file.unlink()
+    raise DownloadError(
+        f"Failed to download {file_name!r} from {download_url!r}"
+    ) from e
+
+
+def _ensure_downloaded_binary(abi: str, tool_name: str) -> pathlib.Path:
+  """Downloads the pre-built binary for the given ABI if not cached."""
+  return _ensure_downloaded_file(
+      abi,
+      tool_name,
+      constants.LITERT_BINARIES_BASE_URL_ANDROID,
+      make_executable=True,
+  )
+
+
+def find_android_binary(tool_name: str, abi: str) -> pathlib.Path:
+  """Locates or downloads an Android executable binary.
+
+  Always downloads from a fixed URL (cached locally).
+
+  Args:
+    tool_name: Binary name (e.g. 'run_model').
+    abi: Target Android CPU ABI.
+
+  Returns:
+    The absolute path to the binary.
+
+  Raises:
+    click.ClickException: If the binary could not be found or downloaded.
+  """
+  try:
+    downloaded_bin = _ensure_downloaded_binary(abi, tool_name)
+    if downloaded_bin.exists():
+      return downloaded_bin
+    else:
+      raise click.ClickException(
+          f"Downloaded binary {tool_name!r} for ABI {abi!r} does not exist."
+      )
+  except DownloadError as e:
+    raise click.ClickException(
+        f"Could not find or download {tool_name!r} for ABI {abi!r}: {e}"
+    ) from e
+
+
+def _ensure_downloaded_library(
+    abi: str,
+    lib_name: str,
+    base_url: str = constants.LITERT_BINARIES_BASE_URL_ANDROID,
+) -> pathlib.Path:
+  """Downloads the pre-built library for the given ABI if not cached."""
+  return _ensure_downloaded_file(abi, lib_name, base_url, make_executable=False)
+
+
+def find_android_lib(
+    lib_name: str,
+    abi: str,
+    base_url: str = constants.LITERT_BINARIES_BASE_URL_ANDROID,
+) -> pathlib.Path:
+  """Locates or downloads an Android library.
+
+  Always downloads from a fixed URL (cached locally).
+
+  Args:
+    lib_name: library name (e.g. 'libLiteRt.so').
+    abi: Target Android CPU ABI.
+    base_url: The base URL to download the library from.
+
+  Returns:
+    The absolute path to the library.
+
+  Raises:
+    click.ClickException: If the library could not be found or downloaded.
+  """
+  try:
+    downloaded_bin = _ensure_downloaded_library(abi, lib_name, base_url)
+    if downloaded_bin.exists():
+      return downloaded_bin
+    else:
+      raise click.ClickException(
+          f"Downloaded library {lib_name!r} for ABI {abi!r} does not exist."
+      )
+  except DownloadError as e:
+    raise click.ClickException(
+        f"Could not find or download {lib_name!r} for ABI {abi!r}: {e}"
+    ) from e
+
+
+def find_npu_dispatch_lib(soc_vendor: str, abi: str) -> pathlib.Path:
+  """Finds and downloads the NPU dispatch library for the given SoC vendor.
+
+  Args:
+    soc_vendor: The NPU vendor ("qualcomm" or "mediatek").
+    abi: Target Android CPU ABI.
+
+  Returns:
+    The absolute path to the local downloaded backend dispatch library.
+  """
+  if soc_vendor == "qualcomm":
+    lib_name = "libLiteRtDispatch_Qualcomm.so"
+  elif soc_vendor == "mediatek":
+    lib_name = "libLiteRtDispatch_MediaTek.so"
+  else:
+    raise click.ClickException(
+        f"Unsupported NPU vendor for dispatch: {soc_vendor!r}"
+    )
+
+  return find_android_lib(
+      lib_name, abi, base_url=constants.LITERT_BINARIES_BASE_URL_ANDROID
+  )
+
+
+def find_npu_compiler_plugin_lib(soc_vendor: str, abi: str) -> pathlib.Path:
+  """Finds and downloads the NPU compiler plugin library for the given SoC vendor.
+
+  Args:
+    soc_vendor: The NPU vendor ("qualcomm" or "mediatek").
+    abi: Target Android CPU ABI.
+
+  Returns:
+    The absolute path to the local downloaded backend compiler plugin library.
+  """
+  if soc_vendor == "qualcomm":
+    lib_name = "libLiteRtCompilerPlugin_Qualcomm.so"
+  elif soc_vendor == "mediatek":
+    lib_name = "libLiteRtCompilerPlugin_MediaTek.so"
+  else:
+    raise click.ClickException(
+        f"Unsupported NPU vendor for compiler plugin: {soc_vendor!r}"
+    )
+
+  return find_android_lib(
+      lib_name, abi, base_url=constants.LITERT_BINARIES_BASE_URL_ANDROID
+  )