chore(model-driver-mediapipe): streamline asset prep with retries (#902)

---------

Co-authored-by: Neko <neko@ayaka.moe>

packages/model-driver-mediapipe/package.json

@@ -24,6 +24,7 @@
   },
   "dependencies": {
     "@mediapipe/tasks-vision": "^0.10.0",
+    "@moeru/std": "catalog:",
     "@pixiv/three-vrm": "^3.4.4",
     "es-toolkit": "catalog:",
     "three": "^0.182.0"

packages/model-driver-mediapipe/tasks/prepare-tasks.ts

@@ -8,6 +8,8 @@ import fs from 'node:fs/promises'
 import { Buffer } from 'node:buffer'
 import { fileURLToPath } from 'node:url'
 
+import { withRetry } from '@moeru/std'
+import { attemptAsync } from 'es-toolkit'
 import { ofetch } from 'ofetch'
 
 import { visionTaskAssets } from './tasks'
@@ -18,39 +20,79 @@ const taskSources: Record<keyof VisionTaskAssets, string> = {
   face: 'https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task',
 }
 
-await fs.mkdir(fileURLToPath(new URL('./assets', import.meta.url)), { recursive: true })
-
-await Promise.all(Object.entries(taskSources).map(
-  async ([key, url]) => {
-    console.log(`Downloading MediaPipe vision task asset for ${key} from ${url}...`)
-    const res = await ofetch(url, { responseType: 'arrayBuffer' })
-    const outputPath = fileURLToPath(visionTaskAssets[key as keyof VisionTaskAssets])
-    await fs.writeFile(outputPath, Buffer.from(res))
-    console.log(`MediaPipe vision task asset for ${key} saved to ${outputPath}`)
-  },
-))
-
+const assetsRoot = fileURLToPath(new URL('./assets', import.meta.url))
 const wasmSourceDir = fileURLToPath(new URL('../node_modules/@mediapipe/tasks-vision/wasm', import.meta.url))
 const wasmOutputDir = fileURLToPath(new URL('./assets/wasm', import.meta.url))
+const taskTargets = Object.entries(taskSources).map(([key, source]) => ({
+  key: key as keyof VisionTaskAssets,
+  source,
+  outputPath: fileURLToPath(visionTaskAssets[key as keyof VisionTaskAssets]),
+}))
+
+async function isUsableFile(path: string) {
+  try {
+    const stat = await fs.stat(path)
+    return stat.isFile() && stat.size > 0
+  }
+  catch {
+    return false
+  }
+}
+
+async function downloadAsset(key: string, url: string, outputPath: string) {
+  const tempPath = `${outputPath}.download`
+  let attempt = 0
+
+  const downloadWithRetry = withRetry(async () => {
+    attempt += 1
+    console.log(`Downloading MediaPipe vision task asset for ${key} from ${url} (attempt ${attempt})...`)
+
+    try {
+      const [fetchError, response] = await attemptAsync(() => ofetch(url, { responseType: 'arrayBuffer' }))
+
+      if (fetchError || !response)
+        throw fetchError ?? new Error(`Missing response while downloading MediaPipe vision task asset for ${key}`)
+
+      await fs.writeFile(tempPath, Buffer.from(response))
+      await fs.rename(tempPath, outputPath)
+      console.log(`MediaPipe vision task asset for ${key} saved to ${outputPath}`)
+    }
+    finally {
+      await fs.rm(tempPath, { force: true })
+    }
+  }, {
+    onError: (error) => {
+      const message = error instanceof Error ? error.message : String(error)
+      console.warn(`Failed to download MediaPipe vision task asset for ${key} (attempt ${attempt}): ${message}`)
+    },
+  })
+
+  try {
+    await downloadWithRetry()
+  }
+  catch (error) {
+    throw new Error(`Failed to download MediaPipe vision task asset for ${key} after ${attempt} attempts`, {
+      cause: error,
+    })
+  }
+}
+
+await fs.mkdir(assetsRoot, { recursive: true })
+
+for (const { key, source, outputPath } of taskTargets) {
+  if (await isUsableFile(outputPath)) {
+    console.log(`MediaPipe vision task asset for ${key} already exists at ${outputPath}, skipping download.`)
+    continue
+  }
+  await downloadAsset(key, source, outputPath)
+
+  if (!await isUsableFile(outputPath))
+    throw new Error(`Failed to ensure MediaPipe vision task asset for ${key}: missing or empty file at ${outputPath}`)
+}
+
 await fs.mkdir(wasmOutputDir, { recursive: true })
 await fs.cp(wasmSourceDir, wasmOutputDir, { recursive: true, force: true })
 
-await Promise.all(Object.entries(visionTaskAssets).map(
-  async ([key, url]) => {
-    const path = fileURLToPath(url)
-    try {
-      await fs.access(path, fs.constants.R_OK)
-    }
-    catch (err) {
-      throw new Error(`Failed to ensure MediaPipe vision task asset for ${key}: ${err}`)
-    }
-    const stat = await fs.stat(path)
-    if (!stat.isFile()) {
-      throw new Error(`Failed to ensure MediaPipe vision task asset for ${key}: not a file: ${path}`)
-    }
-  },
-))
-
 const wasmEntries = await fs.readdir(wasmOutputDir)
 if (!wasmEntries.length)
   throw new Error(`Failed to ensure MediaPipe WASM assets: ${wasmOutputDir} is empty`)

pnpm-lock.yaml

@@ -1808,6 +1808,9 @@ importers:
       '@mediapipe/tasks-vision':
         specifier: ^0.10.0
         version: 0.10.21
+      '@moeru/std':
+        specifier: 'catalog:'
+        version: 0.1.0-beta.14
       '@pixiv/three-vrm':
         specifier: ^3.4.4
         version: 3.4.4(three@0.182.0)