Download AI models with the Background Fetch API

if ('serviceWorker' in navigator) {
  window.addEventListener('load', async () => {
    const registration = await navigator.serviceWorker.register('sw.js');
    console.log('Service worker registered for scope', registration.scope);
  });
}

const FETCH_ID = 'gemma-2b';
const MODEL_URL =
  'https://storage.googleapis.com/jmstore/kaggleweb/grader/g-2b-it-gpu-int4.bin';

downloadButton.addEventListener('click', async (event) => {
  // If the model is already downloaded, return it from the cache.
  const modelAlreadyDownloaded = await caches.match(MODEL_URL);
  if (modelAlreadyDownloaded) {
    const modelBlob = await modelAlreadyDownloaded.blob();
    // Do something with the model.
    console.log(modelBlob);
    return;
  }

  // The model still needs to be downloaded.
  // Feature detection and fallback to classic `fetch()`.
  if (!('BackgroundFetchManager' in self)) {
    try {
      const response = await fetch(MODEL_URL);
      if (!response.ok || response.status !== 200) {
        throw new Error(`Download failed ${MODEL_URL}`);
      }
      const modelBlob = await response.blob();
      // Do something with the model.
      console.log(modelBlob);
      return;
    } catch (err) {
      console.error(err);
    }
  }

  // The service worker registration.
  const registration = await navigator.serviceWorker.ready;

  // Check if there's already a background fetch running for the `FETCH_ID`.
  let bgFetch = await registration.backgroundFetch.get(FETCH_ID);

  // If not, start a background fetch.
  if (!bgFetch) {
    bgFetch = await registration.backgroundFetch.fetch(FETCH_ID, MODEL_URL, {
      title: 'Gemma 2B model',
      icons: [
        {
          src: 'icon.png',
          size: '128x128',
          type: 'image/png',
        },
      ],
      downloadTotal: await getResourceSize(MODEL_URL),
    });
  }
});

const getResourceSize = async (url) => {
  try {
    const response = await fetch(url, { method: 'HEAD' });
    if (response.ok) {
      return response.headers.get('Content-Length');
    }
    console.error(`HTTP error: ${response.status}`);
    return 0;
  } catch (error) {
    console.error('Error fetching content size:', error);
    return 0;
  }
};

bgFetch.addEventListener('progress', (e) => {
  // There's no download progress yet.
  if (!bgFetch.downloadTotal) {
    return;
  }
  // Something went wrong.
  if (bgFetch.failureReason) {
    console.error(bgFetch.failureReason);
  }
  if (bgFetch.result === 'success') {
    return;
  }
  // Update the user about progress.
  console.log(`${bgFetch.downloaded} / ${bgFetch.downloadTotal}`);
});

self.addEventListener('backgroundfetchsuccess', (event) => {
  // Get the background fetch registration.
  const bgFetch = event.registration;

  event.waitUntil(
    (async () => {
      // Open a cache named 'downloads'.
      const cache = await caches.open('downloads');
      // Go over all records in the background fetch registration.
      // (In the running example, there's just one record, but this way
      // the code is future-proof.)
      const records = await bgFetch.matchAll();
      // Wait for the response(s) to be ready, then cache it/them.
      const promises = records.map(async (record) => {
        const response = await record.responseReady;
        await cache.put(record.request, response);
      });
      await Promise.all(promises);

      // Update the browser UI.
      event.updateUI({ title: 'Model downloaded' });

      // Inform the clients that the model was downloaded.
      self.clients.matchAll().then((clientList) => {
        for (const client of clientList) {
          client.postMessage({
            message: 'download-complete',
            id: bgFetch.id,
          });
        }
      });
    })(),
  );
});

navigator.serviceWorker.addEventListener('message', async (event) => {
  const cache = await caches.open('downloads');
  const keys = await cache.keys();
  for (const key of keys) {
    const modelBlob = await cache
      .match(key)
      .then((response) => response.blob());
    // Do something with the model.
    console.log(modelBlob);
  }
});

const registration = await navigator.serviceWorker.ready;
const bgFetch = await registration.backgroundFetch.get(FETCH_ID);
if (!bgFetch) {
  return;
}
await bgFetch.abort();