litert-cli 0.1.0__py3-none-any.whl

litert_cli/commands/download.py

@@ -0,0 +1,274 @@
+# 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 downloading LiteRT models from huggingface.
+
+External packages like requests, or huggingface_hub are optional and will be
+auto-installed when the command is invoked.
+"""
+
+from __future__ import annotations
+
+import datetime
+import json
+import pathlib
+import textwrap
+import urllib.parse
+
+import click
+import huggingface_hub
+import huggingface_hub.utils
+import requests
+
+from ..core import constants
+from ..core import deps
+
+
+def _download_from_url(model_path: str, output_dir: str | pathlib.Path) -> None:
+  """Downloads a model directly from a given URL.
+
+  Args:
+    model_path: The direct URL to the model file.
+    output_dir: The directory where the model will be saved.
+  """
+  click.echo(f"Downloading direct URL: {model_path} to {output_dir}")
+
+  parsed_url = urllib.parse.urlparse(model_path)
+  filename = pathlib.Path(parsed_url.path).name
+  if not filename:
+    filename = "downloaded_model.tflite"
+
+  filepath = pathlib.Path(output_dir) / filename
+
+  try:
+    with requests.get(model_path, stream=True) as response:
+      response.raise_for_status()
+
+      total_size = int(response.headers.get("content-length", 0))
+      chunk_size = 8192
+
+      with click.progressbar(
+          length=total_size,
+          label=f"Downloading {filename}",
+          empty_char=" ",
+          fill_char=click.style("#", fg="green"),
+      ) as bar:
+        with open(filepath, "wb") as f:
+          for chunk in response.iter_content(chunk_size=chunk_size):
+            if chunk:
+              f.write(chunk)
+              bar.update(len(chunk))
+
+    click.secho(f"Downloaded to {filepath}", fg="green")
+  except requests.exceptions.RequestException as e:
+    raise click.ClickException(f"Failed to download URL: {e}") from e
+
+
+def _download_from_hf(
+    repo_id: str,
+    output_dir: str | pathlib.Path,
+    file_pattern: str | None = None,
+    token: str | None = None,
+) -> None:
+  """Downloads a model repository or specific files from HuggingFace Hub.
+
+  Args:
+    repo_id: The HuggingFace Hub repository ID.
+    output_dir: The directory where the files will be saved.
+    file_pattern: Optional glob pattern to filter files to download.
+    token: Optional HuggingFace access token.
+  """
+
+  click.echo(f"Downloading from HuggingFace Hub: {repo_id}")
+
+  kwargs = {
+      "local_dir": output_dir,
+  }
+
+  if token:
+    kwargs["token"] = token
+  if file_pattern:
+    kwargs["allow_patterns"] = file_pattern
+
+  try:
+    downloaded_path = huggingface_hub.snapshot_download(
+        repo_id=repo_id, **kwargs
+    )
+    click.secho(f"Successfully downloaded to {downloaded_path}", fg="green")
+  except huggingface_hub.utils.HfHubHTTPError as e:
+    if "401" in str(e):
+      raise click.ClickException(
+          f"Failed to download from HuggingFace: {e}\n"
+          "This model might require a token. Use --token"
+      ) from e
+    raise click.ClickException(
+        f"Failed to download from HuggingFace: {e}"
+    ) from e
+  except Exception as e:  # pylint: disable=broad-exception-caught
+    # Catching broad exception as fallback for other huggingface_hub errors
+    raise click.ClickException(
+        f"An error occurred during HuggingFace download: {e}"
+    ) from e
+
+
+@click.command(
+    "download",
+    help=textwrap.dedent("""\
+    Download models from URL or HuggingFace.
+
+    MODEL_PATH: Direct URL to a file, or a HuggingFace Hub repository ID.
+
+    Examples:
+
+      Download a direct URL to the current directory:
+
+        $ litert download https://example.com/model.tflite
+
+      Download all .tflite files from a HuggingFace repo:
+
+        $ litert download Qwen/Qwen1.5-0.5B-Chat --file "*.tflite"
+
+      Download from a private repository using a token:
+
+        $ litert download MyOrg/PrivateModel --token hf_your_token
+    """),
+)
+@click.argument("model_path")
+@click.option(
+    "--model-ref",
+    help=(
+        "Reference name for the model in centralized cache (e.g."
+        " 'mobilenet:v3')."
+    ),
+)
+@click.option(
+    "--output",
+    help="Specify output directory (defaults to centralized cache).",
+)
+@click.option(
+    "--file",
+    "file_pattern",
+    help="Only download specific files in the HF repository (e.g. '*.tflite')",
+)
+@click.option("--token", help="Private model requiring a Token")
+@click.option(
+    "--hf-id", help="Optional Hugging Face ID to associate with the model."
+)
+@deps.require_extra("download")
+def download_cmd(
+    model_path: str,
+    model_ref: str | None,
+    output: str | None,
+    file_pattern: str | None,
+    token: str | None,
+    hf_id: str | None,
+) -> None:
+  """Downloads models from URL or HuggingFace.
+
+  Args:
+    model_path: Direct URL to a file, or a HuggingFace Hub repository ID.
+    model_ref: Optional reference name for the model.
+    output: Optional custom directory to save files.
+    file_pattern: Glob pattern to filter downloaded files from HF.
+    token: HuggingFace access token for private repositories.
+    hf_id: Optional Hugging Face ID to associate with the model.
+  """
+  # Resolve output directory and handle sub-ref in model_ref
+  if model_ref and ":" in model_ref:
+    ref, resolved_sub_ref = model_ref.split(":", 1)
+  else:
+    ref = model_ref
+    resolved_sub_ref = None
+
+  if output:
+    output_path = pathlib.Path(output)
+  else:
+    # Use centralized cache by default for HF models
+    if not model_path.startswith(("http://", "https://")):
+      # Default to hf_id as ref if not provided
+      ref = ref or model_path.replace("/", "__")
+      # Flatten ref for directory name consistency
+      ref_flat = ref.replace("/", "__") if "/" in ref else ref
+      output_path = pathlib.Path(constants.LITERT_MODELS_CACHE_DIR) / ref_flat
+    else:
+      output_path = pathlib.Path(".")
+
+  output_path.mkdir(parents=True, exist_ok=True)
+
+  if model_path.startswith(("http://", "https://")):
+    _download_from_url(model_path, output_path)
+  else:
+    _download_from_hf(model_path, output_path, file_pattern, token)
+
+    # Save metadata if it's a managed download (non-URL) to the centralized cache
+    if not model_path.startswith(("http://", "https://")) and not output:
+      final_ref = ref or model_path.replace("/", "__")
+
+      # Load existing metadata if it exists (to preserve other sub-refs)
+      metadata_file = output_path / "metadata.json"
+      metadata = {}
+      try:
+        with open(metadata_file, "r") as f:
+          metadata = json.load(f)
+      except (IOError, json.JSONDecodeError):
+        # If the file doesn't exist or is malformed, start with empty metadata.
+        pass
+
+      # Update or create metadata
+      metadata.setdefault("files", [])
+      metadata.setdefault("sub_references", {})
+
+      metadata.update({
+          "model_ref": final_ref,
+          "hf_id": hf_id or model_path,
+          "source": "huggingface",
+          "created_at": (
+              datetime.datetime.now(datetime.timezone.utc).isoformat()
+          ),
+      })
+
+      if file_pattern and file_pattern not in metadata["files"]:
+        metadata["files"].append(file_pattern)
+
+      # Handle sub-references
+      if resolved_sub_ref:
+        sub_references = metadata["sub_references"]
+
+        # If we downloaded a specific file, find it in the directory!
+        # Since file_pattern might be a glob, we search for matching files.
+        matched_files = []
+        if file_pattern:
+          # Glob search in output_path
+          matched_files = list(output_path.glob(file_pattern))
+
+        file_to_map = ""
+        if matched_files:
+          # Use the first matched file for now
+          file_to_map = matched_files[0].name
+        elif file_pattern and "*" not in file_pattern:
+          # If it's not a glob, assume it's the file name
+          file_to_map = file_pattern
+
+        sub_references[resolved_sub_ref] = {
+            "file": file_to_map,
+            "hf_id": hf_id or model_path,
+        }
+        click.echo(
+            f"Mapped sub-reference '{resolved_sub_ref}' to file '{file_to_map}'"
+        )
+
+      with open(metadata_file, "w") as f:
+        json.dump(metadata, f, indent=2)
+      click.secho(f"Metadata saved to {metadata_file}", fg="green")

litert_cli/commands/import.py

@@ -0,0 +1,124 @@
+# 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 importing local models into LiteRT cache."""
+
+from __future__ import annotations
+
+import datetime
+import json
+import pathlib
+import shutil
+
+import click
+
+from ..core import constants
+
+
+@click.command(
+    "import",
+    help="Import a local file or directory into the centralized cache.",
+)
+@click.argument(
+    "file_path",
+    type=click.Path(
+        exists=True, file_okay=True, dir_okay=True, path_type=pathlib.Path
+    ),
+)
+@click.option(
+    "--model-ref",
+    required=True,
+    help="Reference name for the model (e.g. 'my_model:enc').",
+)
+@click.option("--hf-id", help="Hugging Face ID to associate with the model.")
+def import_cmd(
+    file_path: pathlib.Path, model_ref: str, hf_id: str | None
+) -> None:
+  """Imports a local file or directory into the centralized cache.
+
+  Args:
+    file_path: Path to the local file or directory to import.
+    model_ref: Reference name for management.
+    hf_id: Optional Hugging Face ID.
+  """
+  # Parse model_ref for sub-reference if using ':' syntax
+  main_ref = model_ref
+  resolved_sub_ref = None
+
+  if ":" in model_ref:
+    parts = model_ref.split(":", 1)
+    main_ref = parts[0]
+    resolved_sub_ref = parts[1]
+
+  # Flatten main_ref for directory name consistency (e.g. a/b -> a__b)
+  main_ref_flat = main_ref.replace("/", "__") if "/" in main_ref else main_ref
+  output_path = pathlib.Path(constants.LITERT_MODELS_CACHE_DIR) / main_ref_flat
+  output_path.mkdir(parents=True, exist_ok=True)
+
+  click.echo(f"Importing {file_path} to {output_path}...")
+
+  try:
+    if file_path.is_dir():
+      # Copy directory content, not the directory itself, to avoid nested
+      # folders
+      for item in file_path.iterdir():
+        if item.is_dir():
+          shutil.copytree(item, output_path / item.name, dirs_exist_ok=True)
+        else:
+          shutil.copy(item, output_path / item.name)
+    else:
+      shutil.copy(file_path, output_path / file_path.name)
+
+    click.secho(f"Successfully imported to {output_path}", fg="green")
+
+    # Save metadata
+    metadata_file = output_path / "metadata.json"
+    metadata = {}
+    if metadata_file.exists():
+      try:
+        with open(metadata_file, "r") as f:
+          metadata = json.load(f)
+      except (IOError, json.JSONDecodeError):
+        pass
+
+    metadata.update({
+        "model_ref": main_ref,
+        "hf_id": hf_id,
+        "source": "imported",
+        "import_path": str(file_path.resolve()),
+        "created_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
+    })
+
+    # Handle sub-references
+    if resolved_sub_ref:
+      if "sub_references" not in metadata:
+        metadata["sub_references"] = {}
+
+      # Use the file name in the destination directory
+      file_to_map = file_path.name
+      metadata["sub_references"][resolved_sub_ref] = {
+          "file": file_to_map,
+          "hf_id": hf_id,
+      }
+      click.echo(
+          f"Mapped sub-reference '{resolved_sub_ref}' to file '{file_to_map}'"
+      )
+
+    with open(metadata_file, "w") as f:
+      json.dump(metadata, f, indent=2)
+    click.secho(f"Metadata saved to {metadata_file}", fg="green")
+
+  except Exception as e:
+    raise click.ClickException(f"Failed to import model: {e}") from e

litert_cli/commands/list.py

@@ -0,0 +1,132 @@
+# 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 listing managed models in LiteRT cache."""
+
+from __future__ import annotations
+
+import collections
+import json
+import pathlib
+import sys
+
+import click
+
+from ..core.constants import LITERT_MODELS_CACHE_DIR
+
+
+@click.command(
+    "list",
+    help="List all managed models or detailed contents of a specific model.",
+)
+@click.argument("model_ref", required=False)
+def list_cmd(model_ref: str | None) -> None:
+  """Lists managed models. If MODEL_REF is provided, shows detailed contents."""
+  cache_dir = pathlib.Path(LITERT_MODELS_CACHE_DIR)
+
+  if not cache_dir.exists() or not cache_dir.is_dir():
+    click.echo("No managed models found (cache directory does not exist).")
+    return
+
+  # Case 1: List detailed contents of a specific model
+  if model_ref:
+    # Flatten for directory check
+    ref_flat = model_ref.replace("/", "__") if "/" in model_ref else model_ref
+    model_dir = cache_dir / ref_flat
+
+    if not model_dir.exists() or not model_dir.is_dir():
+      click.secho(f"Error: Managed model '{model_ref}' not found.", fg="red")
+      sys.exit(1)
+
+    metadata_file = model_dir / "metadata.json"
+    metadata = {}
+    if metadata_file.exists():
+      try:
+        with open(metadata_file, "r") as f:
+          metadata = json.load(f)
+      except Exception:
+        pass
+
+    click.echo(
+        "Details for managed model:"
+        f" {click.style(model_ref, fg='green', bold=True)}"
+    )
+    if metadata:
+      click.echo(f"  HF ID:      {metadata.get('hf_id', 'N/A')}")
+      click.echo(f"  Source:     {metadata.get('source', 'N/A')}")
+      click.echo(f"  Created At: {metadata.get('created_at', 'N/A')}")
+
+    click.echo("\nFiles in model directory:")
+    sub_refs = metadata.get("sub_references", {})
+    # Reverse mapping for display
+    file_to_subrefs = collections.defaultdict(list)
+    for sub, info in sub_refs.items():
+      f = info.get("file")
+      if f:
+        file_to_subrefs[f].append(sub)
+
+    for item in sorted(model_dir.iterdir()):
+      if item.name == "metadata.json":
+        continue
+
+      size_kb = item.stat().st_size / 1024
+      suffix = ""
+      if item.name in file_to_subrefs:
+        subs = ", ".join(file_to_subrefs[item.name])
+        suffix = f" {click.style(f'[{subs}]', fg='cyan')}"
+
+      click.echo(f"  - {item.name:<30} ({size_kb:>8.1f} KB){suffix}")
+    return
+
+  # Case 2: List all managed models (Default)
+  models = [d for d in cache_dir.iterdir() if d.is_dir()]
+
+  if not models:
+    click.echo("No managed models found in cache.")
+    return
+
+  click.echo(f"Managed models in {cache_dir}:")
+  click.echo("-" * 60)
+
+  for model_dir in sorted(models, key=lambda x: x.name):
+    ref = model_dir.name
+    metadata_file = model_dir / "metadata.json"
+    hf_id = "N/A"
+    source = "N/A"
+    sub_refs = {}
+
+    if metadata_file.exists():
+      try:
+        with open(metadata_file, "r") as f:
+          metadata = json.load(f)
+          # Use original model_ref from metadata if available
+          ref = metadata.get("model_ref", ref)
+          hf_id = metadata.get("hf_id", "N/A")
+          source = metadata.get("source", "N/A")
+          sub_refs = metadata.get("sub_references", {})
+      except Exception:
+        pass
+
+    click.echo(f"Ref: {click.style(ref, fg='green', bold=True)}")
+    click.echo(f"  HF ID:  {hf_id}")
+    click.echo(f"  Source: {source}")
+
+    if sub_refs:
+      click.echo("  Sub-references:")
+      for sub_ref, info in sub_refs.items():
+        file_name = info.get("file", "N/A")
+        click.echo(f"    - {sub_ref} -> {file_name}")
+
+    click.echo("-" * 60)

litert_cli/commands/lm.py

@@ -0,0 +1,74 @@
+# 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 to call LiteRT-LM CLI tool."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+import click
+from litert_cli.core import deps
+
+
+@click.command(
+    name="lm",
+    context_settings=dict(
+        ignore_unknown_options=True,
+        allow_extra_args=True,
+        help_option_names=[],
+    ),
+    help="""LiteRT-LM CLI commands.
+
+This command is a transparent proxy to the native `litert-lm-cli` package.
+Any wildcard arguments used here are forwarded directly to the actual engine.
+
+Examples:
+
+  Display the native litert-lm help:
+
+    $ litert lm --help
+
+  Run a generative LLM model:
+
+    $ litert lm run model.litertlm
+    $ litert lm run gemma3-1b
+
+  Benchmark a generative LLM model:
+
+    $ litert lm benchmark gemma3-1b
+""",
+)
+@deps.require_extra("lm")
+@click.pass_context
+def lm_cmd(ctx: click.Context) -> None:
+  """LiteRT-LM related commands.
+
+  This command is a transparent proxy to the native `litert-lm-cli` package.
+  Any wildcard arguments used here are forwarded directly to the actual engine.
+
+  Args:
+    ctx: click Context object containing forwarded arguments.
+  """
+  try:
+    result = subprocess.run(["litert-lm"] + ctx.args, check=False)
+    sys.exit(result.returncode)
+  except FileNotFoundError:
+    click.secho("Error: 'litert-lm' executable not found in PATH.", fg="red")
+    click.secho(
+        "Please install 'litert-lm-cli' to use this command.", fg="yellow"
+    )
+    sys.exit(1)