Merge pull request #7 from google-ai-edge:fix-quantize-recipe-bug

PiperOrigin-RevId: 912611645

Author: copybara-github

README.md

@@ -128,14 +128,16 @@ litert compile model.tflite --target sm8750 --target mt6989 --export-aipack my_n
 litert run model.tflite --desktop --cpu
 litert run my_model_ref --desktop --cpu
 
-# Run with GPU acceleration
-litert run model.tflite --gpu
+# Run with GPU acceleration and CPU fallback (multi-accelerator)
+litert run model.tflite --gpu --cpu
+litert run model.tflite --accelerator gpu,cpu
 
 # Run on connected Android device
 litert run model.tflite --android
 
-# Run on connected Android device with NPU acceleration (JIT mode)
-litert run model.tflite --android --npu
+# Run on connected Android device with NPU acceleration and CPU fallback
+litert run model.tflite --android --npu --cpu
+litert run model.tflite --android --accelerator npu,cpu
 
 # Run on connected Android device with NPU AOT-compiled model
 litert run model_sm8450.tflite --android --npu

litert_cli/commands/run/android.py

@@ -22,23 +22,28 @@
   1. Run a model on an Android device:
      $ litert run /path/to/model.tflite --android
 
-  2. Run with custom inputs:
+  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
 
-  3. Run with multiple inputs:
+  4. Run with multiple inputs:
      $ litert run /path/to/model.tflite --android --input input1=value1 --input
      input2=value2
 
-  4. Run with specific signature:
+  5. Run with specific signature:
      $ litert run /path/to/model.tflite --android --signature_index 0
 
-  5. Run with multiple iterations:
+  6. Run with multiple iterations:
      $ litert run /path/to/model.tflite --android --iterations 10
 
-  6. Print tensor details:
+  7. Print tensor details:
      $ litert run /path/to/model.tflite --android --print-tensors
 
-  7. Run with sample size:
+  8. Run with sample size:
      $ litert run /path/to/model.tflite --android --sample-size 100
 """
 
@@ -183,6 +188,7 @@ def run_android(
   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()
 
@@ -270,11 +276,11 @@ def run_android(
   # Pass None as device_id to use the default connected device.
   remote_dispatch_dir = (
       npu.push_npu_runtime_libraries(None, android_root)
-      if accelerator == "npu"
+      if "npu" in accel_list
       else ""
   )
 
-  if accelerator == "npu":
+  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"

litert_cli/commands/run/cli.py

@@ -78,6 +78,14 @@
           7. Print detailed tensor outputs:
 
             $ litert run model.tflite --print-tensors --sample-size 10
+
+          8. Run with multiple accelerators (npu -> gpu -> cpu fallback):
+
+            $ litert run model.tflite --npu --gpu --cpu
+
+            OR explicitly:
+
+            $ litert run model.tflite --accelerator npu,gpu,cpu
         """),
 )
 @deps.require_extra("run")
@@ -118,23 +126,24 @@
     flag_value="android",
     help="Target Android platform to run.",
 )
+@click.option(
+    "--accelerator",
+    type=str,
+    help="Comma-separated list of hardware accelerators (e.g. npu,gpu,cpu).",
+)
 @click.option(
     "--cpu",
-    "accelerator",
-    flag_value="cpu",
-    default=True,
-    help="Use CPU accelerator (Default).",
+    is_flag=True,
+    help="Use CPU accelerator.",
 )
 @click.option(
     "--gpu",
-    "accelerator",
-    flag_value="gpu",
+    is_flag=True,
     help="Use GPU accelerator.",
 )
 @click.option(
     "--npu",
-    "accelerator",
-    flag_value="npu",
+    is_flag=True,
     help="Use NPU accelerator.",
 )
 @click.option(
@@ -169,7 +178,10 @@ def run_cmd(
     model_params: Sequence[str],
     model_help: bool,
     target: str,
-    accelerator: str,
+    accelerator: str | None,
+    cpu: bool,
+    gpu: bool,
+    npu: bool,
     signature_index: int,
     iterations: int,
     print_tensors: bool,
@@ -185,11 +197,33 @@ def run_cmd(
     model_help: Show help specific to the matched model plugin.
     target: Execution target ('desktop' or 'android').
     accelerator: Hardware accelerator ('cpu', 'gpu', or 'npu').
+    cpu: Use CPU accelerator.
+    gpu: Use GPU accelerator.
+    npu: Use NPU accelerator.
     signature_index: Index of model signature to run.
     iterations: Number of times to execute the model for benchmarking.
     print_tensors: Whether to print output tensor elements.
     sample_size: Number of sample elements to print from tensors.
   """
+  # Resolve the order of accelerators
+  accelerator_list = []
+  if accelerator:
+    accelerator_list = [
+        a.strip().lower() for a in accelerator.split(",") if a.strip()
+    ]
+  else:
+    if npu:
+      accelerator_list.append("npu")
+    if gpu:
+      accelerator_list.append("gpu")
+    if cpu:
+      accelerator_list.append("cpu")
+
+    if not accelerator_list:
+      accelerator_list = ["cpu"]
+
+  accelerator = ",".join(accelerator_list)
+
   # Quiet if default is true
   if constants.DEFAULT_QUIET:
 

litert_cli/commands/run/desktop.py

@@ -26,13 +26,18 @@
      OR
      $ litert run /path/to/model.tflite --desktop --accelerator gpu
 
-  3. Run with custom inputs:
+  3. Run with multiple accelerators (gpu -> cpu native fallback):
+     $ litert run /path/to/model.tflite --desktop --gpu --cpu
+     OR
+     $ litert run /path/to/model.tflite --desktop --accelerator gpu,cpu
+
+  4. Run with custom inputs:
      $ litert run /path/to/model.tflite --desktop --input input_name=value
 
-  4. Run with multiple iterations (benchmark):
+  5. Run with multiple iterations (benchmark):
      $ litert run /path/to/model.tflite --desktop --iterations 10
 
-  5. Print tensor details:
+  6. Print tensor details:
      $ litert run /path/to/model.tflite --desktop --print-tensors
 """
 
@@ -240,22 +245,31 @@ def run_desktop(
     click.ClickException: On loading failure or inference execution errors.
   """
 
-  click.echo(
-      f"Loading model on desktop: {model_path} with {accelerator.upper()}"
-  )
+  accel_list = [a.strip().lower() for a in accelerator.split(",") if a.strip()]
 
   # pylint: disable=g-import-not-at-top,reimported
   from ai_edge_litert.compiled_model import CompiledModel
   from ai_edge_litert.compiled_model import Environment
   from ai_edge_litert.hardware_accelerator import HardwareAccelerator
 
-  hw_accel = HardwareAccelerator.CPU
-  if accelerator == "gpu":
-    hw_accel = HardwareAccelerator.GPU
-  elif accelerator == "npu":
-    raise click.ClickException(
-        "NPU accelerator is not yet formally supported via desktop API."
-    )
+  hw_accel = HardwareAccelerator(0)
+  for accel in accel_list:
+    if accel == "cpu":
+      hw_accel |= HardwareAccelerator.CPU
+    elif accel == "gpu":
+      hw_accel |= HardwareAccelerator.GPU
+    elif accel == "npu":
+      hw_accel |= HardwareAccelerator.NPU
+    else:
+      raise click.ClickException(f"Unsupported hardware accelerator: {accel!r}")
+
+  if hw_accel == HardwareAccelerator(0):
+    hw_accel = HardwareAccelerator.CPU
+
+  click.echo(
+      f"Loading model on desktop: {model_path} with native hardware"
+      f" accelerators: {hw_accel}"
+  )
 
   ctx = utils.silence_stderr() if quiet else contextlib.nullcontext()
   with ctx:

test_scripts/models/efficientnet.sh

@@ -53,9 +53,6 @@ export TEST_DATA_DIR="$REPO_ROOT/litert_cli/test_data"
 echo -e "${YELLOW}Installing litert-cli from source...${NC}"
 pip install -e "$REPO_ROOT"
 
-
-
-
 # --- 1. Download EfficientNet-B1 model ---
 run_case "Download: EfficientNet-B1 from HuggingFace" \
     litert download litert-community/efficientnet_b1 --file "*.tflite" --output "$MODEL_DIR/efficientnet"
@@ -113,14 +110,14 @@ if has_android_device; then
     run_case "Benchmark: EfficientNet Dynamic INT8 on Android" \
         litert benchmark "$MODEL_DIR/efficientnet/efficientnet_b1_int8_dynamic.tflite" --android
 else
-    echo -e "\n${YELLOW}No Android device detected. Skipping benchmarks (litert benchmark only supports Android/GCP).${NC}"
+    echo -e "\n${YELLOW}No Android device detected. Skipping benchmarks on Android.${NC}"
 fi
 
 
 # --- 5. Compile (AOT Compilation) ---
 # TODO: Add this back when we fix the NPU compile issue.
 # run_case "Compile: EfficientNet FP32 for Qualcomm sm8750 NPU" \
-#     litert compile "$EFFICIENTNET_TFLITE" --target sm8750 --output-dir "$MODEL_DIR/efficientnet"
+#    litert compile "$EFFICIENTNET_TFLITE" --target sm8750 --output-dir "$MODEL_DIR/efficientnet"
 
 # --- 6. Visualize (Model Explorer) ---
 run_case "Visualize: Launch Model Explorer in the background" \

test_scripts/models/yamnet.sh

@@ -93,9 +93,10 @@ if has_android_device; then
     echo -e "\n${GREEN}Android device detected. Running Android inference...${NC}"
     run_case "Run: YamNet FP32 on Android (CPU)" \
         litert run "$YAMNET_TFLITE" --android --cpu --iterations 1
-
-    run_case "Run: YamNet FP32 on Android (GPU)" \
-        litert run "$YAMNET_TFLITE" --android --gpu --iterations 1
+    
+    # Works on Qualcomm NPU SM8750, but not GPU.
+    # run_case "Run: YamNet FP32 on Android (GPU)" \
+    #    litert run "$YAMNET_TFLITE" --android --gpu --iterations 1
 
     run_case "Run: YamNet Dynamic INT8 on Android (CPU)" \
         litert run "$MODEL_DIR/yamnet/yamnet_int8_dynamic.tflite" --android --cpu --iterations 1
@@ -107,8 +108,9 @@ if has_android_device; then
     run_case "Benchmark: YamNet FP32 on Android (CPU)" \
         litert benchmark "$YAMNET_TFLITE" --android
 
-    run_case "Benchmark: YamNet FP32 on Android (GPU)" \
-        litert benchmark "$YAMNET_TFLITE" --android --gpu
+    # Works on Qualcomm NPU SM8750, but not GPU.
+    # run_case "Benchmark: YamNet FP32 on Android (GPU)" \
+    #    litert benchmark "$YAMNET_TFLITE" --android --gpu
 
     run_case "Benchmark: YamNet Dynamic INT8 on Android" \
         litert benchmark "$MODEL_DIR/yamnet/yamnet_int8_dynamic.tflite" --android