code: def fake_palettize(module, nbits, in_ngroups=1, out_ngroups=1):
      """Simulate weight palettization"""
      from coremltools.models.neural_network.quantization_utils import (
          _get_kmeans_lookup_table_and_weight,
      )
  
      def compress_kmeans(val, nbits):
          lut, indices = _get_kmeans_lookup_table_and_weight(nbits, val)
          lut = lut.astype(val.dtype)
          indices = indices.astype(np.uint8)
          return lut, indices
  
      dtype = module.weight.data.dtype
      device = module.weight.data.device
      val = module.weight.data.cpu().numpy().astype(np.float16)
  
      if out_ngroups == 1 and in_ngroups == 1:
          lut, indices = compress_kmeans(val=val, nbits=nbits)
          module.weight.data = (
              torch.from_numpy(lut[indices])
              .reshape(val.shape)
              .to(dtype)
          )
  
      elif out_ngroups > 1 and in_ngroups == 1:
          assert val.shape[0] % out_ngroups == 0
          rvals = [
              compress_kmeans(val=chunked_val, nbits=nbits)
              for chunked_val in np.split(val, out_ngroups, axis=0)
          ]
          shape = list(val.shape)
          shape[0] = shape[0] // out_ngroups
          module.weight.data = (
              torch.cat(
                  [
                      torch.from_numpy(lut[indices]).reshape(shape)
                      for lut, indices in rvals
                  ],
                  dim=0,
              )
              .to(dtype)
              .to(device)
          )
  
      elif in_ngroups > 1 and out_ngroups == 1:
          assert val.shape[1] % in_ngroups == 0
          rvals = [
              compress_kmeans(val=chunked_val, nbits=nbits)
              for chunked_val in np.split(val, in_ngroups, axis=1)
          ]
          shape = list(val.shape)
          shape[1] = shape[1] // in_ngroups
          module.weight.data = (
              torch.cat(
                  [
                      torch.from_numpy(lut[indices]).reshape(shape)
                      for lut, indices in rvals
                  ],
                  dim=1,
              )
              .to(dtype)
              .to(device)
          )
      else:
          raise ValueError(
              f"in_ngroups={in_ngroups} & out_ngroups={out_ngroups} is illegal!!!"
          )
  
      return torch.from_numpy(val).to(dtype)
  
imports: from collections import OrderedDict
  from copy import deepcopy
  from functools import partial
  import argparse
  import gc
  import json
  import logging
  import numpy
  import os
  from PIL import Image
  from python_coreml_stable_diffusion.torch2coreml import compute_psnr, get_pipeline
  import time
  import torch
  import torch.nn
  import requests
  from tqdm import tqdm, trange
  from coremltools.converters.mil.mil import types
label_desc: Use image node_id 2 for calling a global function with example usage: fake_palettize(module, nbits, in_ngroups, out_ngroups) and returns: torch, lut,  indices