code: def main(args):
      # Load Core ML model
      coreml_model = ct.models.MLModel(
          args.mlpackage_path, compute_units=ct.ComputeUnit.CPU_ONLY
      )
      logger.info(f"Loaded {args.mlpackage_path}")
  
      # Keep track of precision stats
      precision_stats = {
          nbits: {"num_tensors": 0, "numel": 0} for nbits in NBITS
      }
  
      # Load palettization recipe
      with open(args.pre_analysis_json_path, "r") as f:
          pre_analysis = json.load(f)
  
      if args.selected_recipe not in list(pre_analysis["recipes"]):
          raise KeyError(
              f"--selected-recipe ({args.selected_recipe}) not found in "
              f"--pre-analysis-json-path ({args.pre_analysis_json_path}). "
              f" Available recipes: {list(pre_analysis['recipes'])}"
          )
  
      recipe = pre_analysis["recipes"][args.selected_recipe]
      assert all(
          nbits in NBITS + [16] for nbits in recipe.values()
      ), f"Some nbits values in the recipe are illegal. Allowed values: {NBITS}"
  
      # Hash tensors to be able to match torch tensor names to mil tensors
      def get_tensor_hash(tensor):
          assert tensor.dtype == np.float16
          return tensor.ravel()[0] + np.prod(tensor.shape)
  
      args.model_version = pre_analysis["model_version"]
      pipe = get_pipeline(args)
      torch_model = pipe.unet
  
      hashed_recipe = {}
      for torch_module_name, nbits in recipe.items():
          tensor = [
              tensor.cpu().numpy().astype(np.float16)
              for name, tensor in torch_model.named_parameters()
              if name == torch_module_name + ".weight"
          ][0]
          hashed_recipe[get_tensor_hash(tensor)] = nbits
  
      del pipe
      gc.collect()
  
      current_nbits: int
  
      def op_selector(const):
          parameter_tensor = const.val.val
          if parameter_tensor.size < MIN_SIZE:
              return False
  
          if parameter_tensor.dtype != np.float16:
              # These are the tensors that were compressed to look-up indices in previous passes
              return False
  
          tensor_hash = get_tensor_hash(parameter_tensor)
          tensor_spec = (
              f"{tensor_hash} with shape {parameter_tensor.shape}"
          )
  
          hashes = list(hashed_recipe)
          pdist = np.abs(np.array(hashes) - tensor_hash)
          matched = pdist.argmin()
          logger.debug(
              f"{tensor_spec}: {tensor_hash} matched with {hashes[matched]} (hash error={pdist.min()})"
          )
  
          target_nbits = hashed_recipe[hashes[matched]]
  
          do_palettize = current_nbits == target_nbits
          if do_palettize:
              logger.debug(
                  f"{tensor_spec}: Palettizing to {target_nbits}-bit palette"
              )
              precision_stats[current_nbits]["num_tensors"] += 1
              precision_stats[current_nbits]["numel"] += np.prod(
                  parameter_tensor.shape
              )
              return True
          return False
  
      for nbits in NBITS:
          logger.info(
              f"Processing tensors targeting {nbits}-bit palettes"
          )
          current_nbits = nbits
  
          config = ct.optimize.coreml.OptimizationConfig(
              global_config=ct.optimize.coreml.OpPalettizerConfig(
                  mode="kmeans",
                  nbits=nbits,
                  weight_threshold=None,
              ),
              is_deprecated=True,
              op_selector=op_selector,
          )
          coreml_model = ct.optimize.coreml.palettize_weights(
              coreml_model, config=config
          )
          logger.info(
              f"{precision_stats[nbits]['num_tensors']} tensors are palettized with {nbits} bits"
          )
  
      tot_numel = sum(
          [precision_stats[nbits]["numel"] for nbits in NBITS]
      )
      final_size = sum(
          [precision_stats[nbits]["numel"] * nbits for nbits in NBITS]
      )
      logger.info(
          f"Palettization result: {final_size / tot_numel:.2f}-bits resulting in {final_size / (8*1e6)} MB"
      )
      pprint(precision_stats)
      coreml_model.save(args.o)
  
file_path: ml-stable-diffusion/python_coreml_stable_diffusion/mixed_bit_compression_apply.py
imports: from pprint import pprint
  import argparse
  import coremltools
  import gc
  import json
  import logging
  import numpy
  import os
  from python_coreml_stable_diffusion.torch2coreml import get_pipeline
  from python_coreml_stable_diffusion.mixed_bit_compression_pre_analysis import NBITS, PALETTIZE_MIN_SIZE
label_desc: Use image node_id 1 for calling a global function with example usage: main(args) and returns: unknown, False, False, False, True