litert-cli 0.1.0__py3-none-any.whl

litert_cli/commands/convert/cli.py

@@ -0,0 +1,255 @@
+# 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.
+# ==============================================================================
+
+"""CLI module for the `litert convert` command.
+
+This command handles converting PyTorch and Hugging Face models into LiteRT
+models using various conversion paths like automated HF export, generic script
+injection, and native generative API re-authoring.
+"""
+
+from __future__ import annotations
+
+import pathlib
+import textwrap
+
+import click
+from litert_cli.core import deps
+
+
+@click.command(
+    "convert",
+    help=textwrap.dedent("""\
+        Convert a PyTorch model into a LiteRT model.
+
+        MODEL_OR_SCRIPT: Hugging Face model ID or path to a PyTorch script.
+
+        Note: Only AutoModelForCausalLM models are supported, when using HF mode.
+
+        Examples:
+
+          Automated HF Conversion:
+
+            $ litert convert Qwen/Qwen1.5-0.5B-Chat --output /tmp/qwen
+
+          HF Conversion with Weight-Only INT4 Quantization:
+
+            $ litert convert Qwen/Qwen1.5-0.5B-Chat --quantize-recipe weight_only_wi4_afp32 --output /tmp/qwen_w4
+
+          HF Conversion with Custom Prefill & Cache Lengths:
+
+            $ litert convert Qwen/Qwen1.5-0.5B-Chat --prefill-lengths "128,512" --cache-length 2048 --output /tmp/qwen_custom
+
+          Generic Script Injection with Quantization:
+
+            $ litert convert my_model.py --quantize-recipe dynamic_wi8_afp32 --output /tmp/mymodel
+
+          Generic Script with Model Args and AOT Target Compilation:
+
+            $ litert convert my_model.py --model-args "batch_size=4" --target sm8450 --output /tmp/mymodel_npu
+    """),
+)
+@deps.require_extra("convert")
+@click.argument("model_or_script", type=str, required=True)
+@click.option(
+    "--output",
+    type=click.Path(
+        file_okay=False,
+        dir_okay=True,
+        resolve_path=True,
+        path_type=pathlib.Path,
+    ),
+    required=False,
+    default=None,
+    help="Directory to save the converted TFLite model.",
+)
+@click.option(
+    "--model-func",
+    type=str,
+    default="get_model",
+    help=(
+        "Name of the function in the --script that returns a torch.nn.Module"
+        " and optionally a quantization config. Default: 'get_model'."
+    ),
+)
+@click.option(
+    "--input-func",
+    type=str,
+    default="get_args",
+    help=(
+        "Name of the function in the --script that returns (True) sample_args "
+        "and/or kwargs. Default: 'get_args'."
+    ),
+)
+@click.option(
+    "--target",
+    type=str,
+    multiple=True,
+    help=(
+        "One or more NPU target codenames (e.g., sm8450) to apply AOT"
+        " compilation."
+    ),
+)
+@click.option(
+    "--export-aipack",
+    type=click.Path(
+        file_okay=False,
+        dir_okay=True,
+        resolve_path=True,
+        path_type=pathlib.Path,
+    ),
+    required=False,
+    default=None,
+    help=(
+        "If specified, exports an AI Pack directory for PODAI alongside the"
+        " compiled model."
+    ),
+)
+@click.option(
+    "--quantize",
+    "--quantize-recipe",
+    type=str,
+    default=None,
+    help=(
+        "Quantization recipe to apply (e.g., dynamic_wi8_afp32,"
+        " weight_only_wi4_afp32). Alias: --quantize-recipe. For full list of"
+        " generative recipes, see 'ai_edge_quantizer.recipe'."
+    ),
+)
+@click.option(
+    "--model-args",
+    type=str,
+    default=None,
+    help=(
+        "Comma-separated key=value arguments to pass to custom model/input"
+        " functions."
+    ),
+)
+@click.option(
+    "--prefill-lengths",
+    type=str,
+    default="256",
+    help=(
+        "Comma-separated list of prefill lengths for HuggingFace models."
+        " Default: '256'."
+    ),
+)
+@click.option(
+    "--cache-length",
+    type=int,
+    default=4096,
+    help="KV cache length for HuggingFace models. Default: 4096.",
+)
+@click.option(
+    "--bundle-litert-lm/--no-bundle-litert-lm",
+    is_flag=True,
+    default=True,
+    help=(
+        "Bundle exported artifacts into a .litert_lm package (HuggingFace mode"
+        " only). Default: True."
+    ),
+)
+def convert_cmd(
+    model_or_script: str,
+    output: pathlib.Path | None,
+    model_func: str,
+    input_func: str,
+    target: tuple[str, ...],
+    export_aipack: pathlib.Path | None,
+    quantize: str | None,
+    model_args: str | None,
+    prefill_lengths: str,
+    cache_length: int,
+    bundle_litert_lm: bool,
+) -> None:
+  r"""Converts a PyTorch model into a LiteRT model.
+
+  Args:
+    model_or_script: Hugging Face model ID or path to a PyTorch script.
+    output: Output directory for the converted model.
+    model_func: Function to retrieve the model in 'script' mode.
+    input_func: Function to retrieve sample inputs in 'script' mode.
+    target: NPU targets to compile for.
+    export_aipack: Output directory to export the AI Pack for PODAI.
+    quantize: Quantization recipe to apply (see ai_edge_quantizer.recipe for
+      full list).
+    model_args: Arguments to pass to custom model/input functions.
+    prefill_lengths: List of prefill lengths for HuggingFace models.
+    cache_length: KV cache length for HuggingFace models.
+    bundle_litert_lm: Whether to bundle artifacts into a .litert_lm package.
+  """
+
+  from litert_cli.core import constants, utils
+  import warnings
+
+  if constants.DEFAULT_QUIET:
+    utils.enable_quiet_mode()
+
+  # Suppress noisy warnings from torch, torchao, etc.
+  warnings.filterwarnings("ignore", category=FutureWarning)
+  warnings.filterwarnings("ignore", category=SyntaxWarning)
+  warnings.filterwarnings("ignore", category=UserWarning)
+
+  if output is None:
+    if model_or_script.endswith(".py"):
+      base_name = pathlib.Path(model_or_script).stem
+    else:
+      base_name = pathlib.Path(model_or_script).name
+    output = pathlib.Path.cwd() / base_name
+
+  if constants.ENABLE_MODEL_PLUGINS:
+    from litert_cli.models import dispatch_model_intent
+
+    plugin_result = dispatch_model_intent(
+        "convert",
+        model_or_script,
+        output=output,
+        target=target,
+        quantize=quantize,
+        export_aipack=export_aipack,
+        model_args=model_args,
+        prefill_lengths=prefill_lengths,
+        cache_length=cache_length,
+        bundle_litert_lm=bundle_litert_lm,
+    )
+    if plugin_result is not None:
+      return
+
+  if model_or_script.endswith(".py"):
+    from litert_cli.commands.convert import generic  # pylint: disable=g-import-not-at-top
+
+    generic.convert_generic_script(
+        model_or_script,
+        model_func,
+        input_func,
+        str(output),
+        target,
+        export_aipack,
+        quantize,
+        model_args,
+    )
+  else:
+    from litert_cli.commands.convert import huggingface  # pylint: disable=g-import-not-at-top
+
+    huggingface.convert_huggingface(
+        model_or_script,
+        str(output),
+        target,
+        export_aipack,
+        quantize,
+        prefill_lengths,
+        cache_length,
+        bundle_litert_lm,
+    )

litert_cli/commands/convert/generic.py

@@ -0,0 +1,211 @@
+# 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.
+# ==============================================================================
+
+"""Generic PyTorch script conversion logic for LiteRT CLI."""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import shutil
+import sys
+from typing import Any
+import uuid
+
+import click
+from litert_cli.core import npu_utils
+
+from ai_edge_litert.aot.ai_pack import export_lib as ai_pack_export
+
+
+def convert_generic_script(
+    script: str,
+    model_func: str,
+    input_func: str,
+    output: str,
+    target: tuple[str, ...],
+    export_aipack: pathlib.Path | None,
+    quantize: str | None = None,
+    model_args: str | None = None,
+) -> None:
+  """Converts using generic PyTorch scripts and `litert_torch.convert`.
+
+  Args:
+    script: Path to the PyTorch script (.py).
+    model_func: Name of function returning the `torch.nn.Module`.
+    input_func: Name of function returning sample inputs.
+    output: Directory to save the converted model.
+    target: NPU targets to apply AOT compilation.
+    export_aipack: Output directory to export the AI Pack for PODAI.
+    quantize: Quantization recipe to apply.
+    model_args: Arguments to pass to custom model/input functions.
+
+  Raises:
+    ImportError: If the script loading fails.
+    AttributeError: If functions are missing in the user script.
+    ValueError: If inputs result shape is not supported.
+  """
+  # pylint: disable=g-import-not-at-top
+  import litert_torch
+
+  script_path = pathlib.Path(script).resolve()
+  click.echo(f"Loading custom script from: {script_path}")
+
+  # Dynamically load the user's python file with sandbox isolation
+  module_name = f"user_model_script_{uuid.uuid4().hex}"
+  spec = importlib.util.spec_from_file_location(module_name, script_path)
+  if spec is None or spec.loader is None:
+    raise ImportError(f"Could not load script {script_path}")
+
+  user_module = importlib.util.module_from_spec(spec)
+  sys.modules[module_name] = user_module
+
+  # Add the script's directory to sys.path so it can resolve its own local imports
+  sys.path.insert(0, str(script_path.parent))
+
+  try:
+    spec.loader.exec_module(user_module)
+
+    # Look up the model factory
+    if not hasattr(user_module, model_func):
+      raise AttributeError(
+          f"Function '{model_func}' not found in {script_path}"
+      )
+    # Look up the args factory
+    if not hasattr(user_module, input_func):
+      raise AttributeError(
+          f"Function '{input_func}' not found in {script_path}"
+      )
+
+    # Parse model_args
+    parsed_args = {}
+    if model_args:
+      for item in model_args.split(","):
+        if "=" in item:
+          k, v = item.split("=", 1)
+          if v.isdigit():
+            parsed_args[k] = int(v)
+          elif v.replace(".", "", 1).isdigit():
+            parsed_args[k] = float(v)
+          elif v.lower() == "true":
+            parsed_args[k] = True
+          elif v.lower() == "false":
+            parsed_args[k] = False
+          else:
+            parsed_args[k] = v
+
+    click.echo(
+        f"Instantiating model via '{model_func}' with args: {parsed_args}..."
+    )
+    model_result = getattr(user_module, model_func)(**parsed_args)
+
+    # The user might return just the nn.Module or a tuple of (nn.Module, QuantConfig)
+    if isinstance(model_result, tuple):
+      model, quant_config = model_result
+    else:
+      model = model_result
+      quant_config = None
+
+    if quant_config is None and quantize:
+      click.echo(f"Building dynamic QuantConfig for recipe '{quantize}'...")
+      from litert_torch.quantize import quant_config as quant_config_lib
+
+      if "pt2e" in quantize.lower():
+        from litert_torch.quantize import pt2e_quantizer as pt2eq
+
+        is_dynamic = "dynamic" in quantize.lower()
+        is_per_channel = "per_channel" in quantize.lower()
+        pt2e_cfg = pt2eq.get_symmetric_quantization_config(
+            is_per_channel=is_per_channel, is_dynamic=is_dynamic
+        )
+        q = pt2eq.PT2EQuantizer().set_global(pt2e_cfg)
+        quant_config = quant_config_lib.QuantConfig(pt2e_quantizer=q)
+      else:
+        from litert_torch.generative.quantize import quant_recipe
+        from litert_torch.generative.quantize import quant_recipe_utils
+
+        if "weight_only" in quantize.lower():
+          recipe_obj = quant_recipe_utils.create_layer_quant_weight_only()
+        elif "fp16" in quantize.lower():
+          recipe_obj = quant_recipe_utils.create_layer_quant_fp16()
+        else:
+          recipe_obj = quant_recipe_utils.create_layer_quant_dynamic()
+
+        quant_config = quant_config_lib.QuantConfig(
+            generative_recipe=quant_recipe.GenerativeQuantRecipe(
+                default=recipe_obj
+            )
+        )
+
+    click.echo(
+        f"Generating sample inputs via '{input_func}' with args:"
+        f" {parsed_args}..."
+    )
+    inputs_result = getattr(user_module, input_func)(**parsed_args)
+
+    sample_args: tuple[Any, ...]
+    sample_kwargs: dict[str, Any]
+    if isinstance(inputs_result, tuple):
+      sample_args = inputs_result
+      sample_kwargs = {}
+    elif isinstance(inputs_result, dict):
+      sample_args = ()
+      sample_kwargs = inputs_result
+    else:
+      raise ValueError(
+          f"'{input_func}' must return a tuple (args) or a dict (kwargs)."
+      )
+
+    click.echo("Executing liteRT conversion tracer...")
+
+    builder = litert_torch
+    if target:
+      for t in target:
+        target_obj = npu_utils.get_target(t)
+        builder = builder.experimental_add_compilation_backend(target_obj)
+
+    edge_model = builder.convert(
+        module=model,
+        sample_args=sample_args,
+        sample_kwargs=sample_kwargs,
+        quant_config=quant_config,
+    )
+
+    out_path = pathlib.Path(output)
+    out_path.mkdir(parents=True, exist_ok=True)
+    model_name = script_path.stem
+
+    if target and export_aipack:
+      export_dir = pathlib.Path(export_aipack)
+      click.echo(f"Exporting AI Pack to: {export_dir}")
+      shutil.rmtree(export_dir, ignore_errors=True)
+      export_dir.mkdir(parents=True, exist_ok=True)
+      ai_pack_export.export(edge_model, str(export_dir), model_name, "model")
+    else:
+      if target:
+        click.echo(f"Exporting compiled model to {out_path} as {model_name}...")
+        edge_model.export(str(out_path), model_name=model_name)
+      else:
+        final_path = out_path / f"{model_name}.tflite"
+        click.echo(f"Exporting converted model to {final_path}...")
+        edge_model.export(str(final_path))
+
+    click.echo("Done!")
+
+  finally:
+    # Cleanup injected sys.path modification and sandbox module
+    sys.modules.pop(module_name, None)
+    if sys.path[0] == str(script_path.parent):
+      sys.path.pop(0)

litert_cli/commands/convert/huggingface.py

@@ -0,0 +1,175 @@
+# 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.
+# ==============================================================================
+
+"""Hugging Face automated export logic for LiteRT conversion.
+
+This module provides the implementation for the 'hf' mode of the
+`litert convert` command, orchestrating the download and conversion of
+models directly from the Hugging Face Hub.
+"""
+
+from __future__ import annotations
+
+import pathlib
+import shutil
+
+import click
+from litert_cli.core import npu_utils
+
+from ai_edge_litert.aot import aot_compile as aot_lib
+from ai_edge_litert.aot.ai_pack import export_lib as ai_pack_export
+
+def convert_huggingface(
+    model: str,
+    output: str,
+    target: tuple[str, ...],
+    export_aipack: pathlib.Path | None,
+    quantize: str | None = None,
+    prefill_lengths: str = "256",
+    cache_length: int = 4096,
+    bundle_litert_lm: bool = True,
+) -> None:
+  """Converts models using HuggingFace Automated Export (export_hf).
+
+  Args:
+    model: The Hugging Face model ID (e.g., Qwen/Qwen1.5-0.5B-Chat).
+    output: The directory to save the converted model.
+    target: NPU targets to apply AOT compilation.
+    export_aipack: Output directory to export the AI Pack for PODAI.
+    quantize: Quantization recipe to apply.
+    prefill_lengths: Comma-separated list of prefill lengths.
+    cache_length: KV cache length.
+    bundle_litert_lm: Whether to bundle artifacts into a .litert_lm package.
+  """
+  # Lazy load the export module to avoid importing torch and other heavy
+  # dependencies when the litert CLI is merely invoked for --help.
+  # pylint: disable=g-import-not-at-top
+  from litert_torch.generative.export_hf import export as hf_export
+  import transformers
+
+  click.echo(f"Starting conversion for model '{model}''")
+
+  try:
+    is_causal_lm = False
+    is_gemma3 = False
+    is_gemma3n = False
+    is_gemma4 = False
+    is_gemma_vlm = False
+
+    # Verify AutoModelForCausalLM architecture
+    try:
+      config = transformers.AutoConfig.from_pretrained(
+          model, trust_remote_code=True
+      )
+      architectures = getattr(config, "architectures", [])
+      is_causal_lm = any("CausalLM" in arch for arch in architectures)
+      is_gemma3 = any(
+          "Gemma3ForConditionalGeneration" in arch for arch in architectures
+      )
+      is_gemma3n = any(
+          "Gemma3nForConditionalGeneration" in arch for arch in architectures
+      )
+      is_gemma4 = any(
+          "Gemma4ForConditionalGeneration" in arch for arch in architectures
+      )
+      is_gemma_vlm = is_gemma3 or is_gemma3n or is_gemma4
+
+      if not (is_causal_lm or is_gemma_vlm):
+        raise ValueError(
+            "Currently only AutoModelForCausalLM is supported (or Gemma VLM"
+            f" architectures: Gemma3, Gemma3n, Gemma4). Model '{model}' has"
+            f" architectures {architectures}."
+        )
+    except Exception as e:
+      if isinstance(e, ValueError):
+        raise
+      click.echo(f"Warning during config verification: {e}", err=True)
+
+    # Parse prefill_lengths
+    parsed_prefill = [int(x.strip()) for x in prefill_lengths.split(",")]
+
+    # Call the auto-export function from litert_torch.
+    # It automatically saves to the output.
+    task = "text_generation"
+    export_kwargs = {}
+    use_jinja_template = is_gemma4
+    if is_gemma_vlm:
+      task = "image_text_to_text"
+      export_kwargs["export_vision_encoder"] = True
+      export_kwargs["externalize_embedder"] = True
+      if is_gemma3 or is_gemma3n:
+        export_kwargs["vision_encoder_quantization_recipe"] = (
+            "weight_only_wi8_afp32"
+        )
+      if is_gemma4:
+        export_kwargs["jinja_chat_template_override"] = (
+            "litert-community/gemma-4-E2B-it-litert-lm"
+        )
+
+    hf_export.export(
+        model=model,
+        output_dir=output,
+        task=task,
+        quantization_recipe=quantize,
+        prefill_lengths=parsed_prefill,
+        cache_length=cache_length,
+        bundle_litert_lm=bundle_litert_lm,
+        trust_remote_code=False,
+        use_jinja_template=use_jinja_template,
+        **export_kwargs,
+    )
+
+    if target:
+      output_dir = pathlib.Path(output)
+      # Find the generated tflite
+      tflite_files = list(output_dir.glob("*.tflite"))
+      if not tflite_files:
+        raise FileNotFoundError(
+            f"No .tflite files found in HF export output: {output_dir}"
+        )
+      target_tflite = tflite_files[0]
+      base_name = target_tflite.stem
+
+      click.echo(
+          f"Compiling converted model {target_tflite} for targets:"
+          f" {', '.join(target)}"
+      )
+      aot_targets = [npu_utils.get_target(t) for t in target]
+
+      compiled_models = aot_lib.aot_compile(
+          str(target_tflite),
+          target=aot_targets,
+          keep_going=False,
+      )
+
+      if export_aipack:
+        export_dir = pathlib.Path(export_aipack)
+        click.echo(f"Exporting AI Pack to: {export_dir}")
+        shutil.rmtree(export_dir, ignore_errors=True)
+        export_dir.mkdir(parents=True, exist_ok=True)
+        ai_pack_export.export(
+            compiled_models, str(export_dir), base_name, "model"
+        )
+      else:
+        # Overwrite the original tflite output with the compiled models
+        click.echo(f"Exporting compiled model over original in {output_dir}")
+        compiled_models.export(str(output_dir), model_name=base_name)
+
+    click.echo(f"Successfully converted and saved model to {output}")
+
+  except Exception as e:
+    click.echo(f"Error during conversion: {e}")
+    # Re-raise to let click handle the error exit
+    raise

litert_cli/commands/delete.py

@@ -0,0 +1,56 @@
+# 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 deleting managed models from LiteRT cache."""
+
+from __future__ import annotations
+
+import pathlib
+import shutil
+
+import click
+
+from ..core.constants import LITERT_MODELS_CACHE_DIR
+
+
+@click.command(
+    "delete", help="Delete a managed model from the centralized cache."
+)
+@click.argument("ref")
+@click.option("--yes", "-y", is_flag=True, help="Do not ask for confirmation.")
+def delete_cmd(ref: str, yes: bool) -> None:
+  """Deletes a managed model from the centralized cache.
+
+  Args:
+    ref: The model reference to delete.
+    yes: Whether to skip confirmation prompt.
+  """
+  # Flatten for directory check
+  ref_flat = ref.replace("/", "__") if "/" in ref else ref
+  cache_path = pathlib.Path(LITERT_MODELS_CACHE_DIR) / ref_flat
+
+  if not cache_path.exists() or not cache_path.is_dir():
+    raise click.ClickException(f"Model reference '{ref}' not found in cache.")
+
+  if not yes:
+    if not click.confirm(f"Are you sure you want to delete model '{ref}'?"):
+      click.echo("Aborted.")
+      return
+
+  try:
+    shutil.rmtree(cache_path)
+    click.secho(f"Successfully deleted model '{ref}' from cache.", fg="green")
+  except OSError as e:
+    raise click.ClickException(f"Failed to delete model: {e}") from e