1import localforage from 'localforage'
  2 
  3export type modelType = 'inpaint' | 'superResolution'
  4 
  5localforage.config({
  6  name: 'modelCache',
  7})
  8 
  9export async function saveModel(modelType: modelType, modelBlob: ArrayBuffer) {
 10  await localforage.setItem(getModel(modelType).name, modelBlob)
 11}
 12 
 13function getModel(modelType: modelType) {
 14  if (modelType === 'inpaint') {
 15    return {
 16      name: 'migan-pipeline-v2',
 17      url: '/models/migan_pipeline_v2.onnx',
 18      backupUrl: '',
 19    }
 20  }
 21  if (modelType === 'superResolution') {
 22    return {
 23      name: 'realesrgan-x4',
 24      url: '/models/realesrgan-x4.onnx',
 25      backupUrl: '',
 26    }
 27  }
 28  throw new Error('wrong modelType')
 29}
 30 
 31export async function loadModel(modelType: modelType): Promise<ArrayBuffer> {
 32  const model = (await localforage.getItem(
 33    getModel(modelType).name
 34  )) as ArrayBuffer
 35  return model
 36}
 37 
 38export async function modelExists(modelType: modelType) {
 39  const model = await loadModel(modelType)
 40  return model !== null && model !== undefined
 41}
 42 
 43export async function ensureModel(modelType: modelType) {
 44  if (await modelExists(modelType)) {
 45    return loadModel(modelType)
 46  }
 47  const model = getModel(modelType)
 48  const response = await fetch(model.url)
 49  const buffer = await response.arrayBuffer()
 50  await saveModel(modelType, buffer)
 51  return buffer
 52}
 53 
 54export async function downloadModel(
 55  modelType: modelType,
 56  setDownloadProgress: (arg0: number) => void
 57) {
 58  if (await modelExists(modelType)) {
 59    return
 60  }
 61 
 62  async function downloadFromUrl(url: string) {
 63    console.log('start download from', url)
 64    setDownloadProgress(0)
 65    const response = await fetch(url)
 66    const fullSize = response.headers.get('content-length')
 67    const reader = response.body!.getReader()
 68    const total: Uint8Array[] = []
 69    let downloaded = 0
 70 
 71    while (true) {
 72      const { done, value } = await reader.read()
 73 
 74      if (done) {
 75        break
 76      }
 77 
 78      downloaded += value?.length || 0
 79 
 80      if (value) {
 81        total.push(value)
 82      }
 83 
 84      setDownloadProgress((downloaded / Number(fullSize)) * 100)
 85    }
 86 
 87    const buffer = new Uint8Array(downloaded)
 88    let offset = 0
 89    for (const chunk of total) {
 90      buffer.set(chunk, offset)
 91      offset += chunk.length
 92    }
 93 
 94    await saveModel(modelType, buffer)
 95    setDownloadProgress(100)
 96  }
 97 
 98  const model = getModel(modelType)
 99  try {
100    await downloadFromUrl(model.url)
101  } catch (e) {
102    if (model.backupUrl) {
103      try {
104        await downloadFromUrl(model.backupUrl)
105      } catch (r) {
106        alert(`Failed to download the backup model: ${r}`)
107      }
108    }
109    alert(`Failed to download the model, network problem: ${e}`)
110  }
111}

 1import { defineConfig } from 'vitest/config'
 2import react from '@vitejs/plugin-react-swc'
 3import eslintPlugin from 'vite-plugin-eslint'
 4 
 5// https://vitejs.dev/config/
 6export default defineConfig({
 7  base: '/',
 8  plugins: [react(), eslintPlugin()],
 9  test: {
10    globals: true,
11    environment: 'jsdom',
12    setupFiles: './src/setupTests.ts',
13    css: true,
14    reporters: ['verbose'],
15    coverage: {
16      reporter: ['text', 'json', 'html'],
17      include: ['src/**/*'],
18      exclude: [],
19    },
20  },
21  server: {
22    headers: {
23      'Cross-Origin-Opener-Policy': 'same-origin',
24      'Cross-Origin-Embedder-Policy': 'credentialless',
25    },
26  },
27})