<!DOCTYPE html>
<html>
<head>
  <script src="face-api.js"></script>
  <script src="FileSaver.js"></script>
  <script src="quantization.js"></script>
</head>
<body>
  <script>
    tf = faceapi.tf

    const uncompressedWeightsUri = `tiny_face_detector_model.weights`
    const net = new faceapi.TinyFaceDetector()

    async function load() {
      await net.load(new Float32Array(await (await fetch(uncompressedWeightsUri)).arrayBuffer()))
      console.log('net loaded')
    }

    function getNamedTensors() {
      return net.getParamList().map(({ path, tensor }) => ({ name: path, tensor }))
    }

    const modelName = 'tiny_face_detector'

    function makeShards(weightArray) {
      const maxLength = 4096 * 1024

      const shards = []

      let shardIdx = 0
      for (let i = 0; i < weightArray.length; i++) {
        if (!shards[shardIdx]) {
          shards[shardIdx] = []
        }

        shards[shardIdx].push(weightArray[i])

        if (shards[shardIdx].length >= maxLength) {
          shardIdx += 1
        }
      }

      return shards.map(shardArray => new Uint8Array(shardArray))
    }

    async function quantizeAndSave() {

      const quantizedTensorArrays = []
      const weightEntries = []

      getNamedTensors().forEach(({ name, tensor, isSkipQuantization }) => {

        const weightEntry = {
          name,
          shape: tensor.shape,
          dtype: tensor.dtype
        }

        if (isSkipQuantization) {
          quantizedTensorArrays.push(new Uint8Array(tensor.dataSync().buffer))
          weightEntries.push(weightEntry)
          return
        }

        const { scale, min, qdata } = quantizeWeights(tensor)

        console.log(name, { scale, min })

        const quantizedWeightEntry = {
          ...weightEntry,
          quantization: { dtype: 'uint8', scale, min }
        }

        quantizedTensorArrays.push(qdata)
        weightEntries.push(quantizedWeightEntry)
      })

      const quantizedWeights = quantizedTensorArrays
        .map(typedArray => Array.from(typedArray))
        .reduce((flat, arr) => flat.concat(arr))

      const shards = makeShards(quantizedWeights)
      console.log('num shards: ', shards.length)

      const paths = []
      shards.forEach((shardData, i) => {
        const shardName = `${modelName}_model-shard${i + 1}`
        paths.push(shardName)
        saveAs(new Blob([shardData]), shardName)
      })

      const weightManifest = [{
        weights: weightEntries,
        paths
      }]

      saveAs(new Blob([JSON.stringify(weightManifest)]), `${modelName}_model-weights_manifest.json`)
    }

    load()

  </script>
</body>
</html>