Files
moeka-project/apps/stage-pocket/vite.config.ts
T
147d077aa7 feat(inference): unify and optimize WebGPU inference pipeline (#1622)
## Problem

The WebGPU inference pipeline had several structural issues:

1. **No unified protocol** — Kokoro TTS, Whisper ASR, and Background
Removal workers each used their own ad-hoc message formats. Adding a new
model meant reinventing worker communication from scratch.
2. **Infrastructure existed but was disconnected** —
`GPUResourceCoordinator`, `LoadQueue`, `InferenceWorkerManager`, and
`protocol.ts` were all implemented but had zero consumers. The adapters
duplicated the same lifecycle/timeout/mutex patterns independently.
3. **Performance gaps** — Kokoro only offered fp32 on WebGPU (no fp16),
Whisper warm-up compiled shaders for 187.5s of dummy audio, audio
transfer went through unnecessary WAV blob encode/decode, and
`listVoices` reloaded the model every time.
4. **Silent failures** — Whisper worker's `generate()` had no try-catch;
errors were swallowed and the main thread waited until timeout.
5. **No graceful degradation** — Whisper and Background Removal workers
hardcoded `device: 'webgpu'` with no WASM fallback.
6. **No observability** — Only Kokoro had performance tracing. No
adapter reported status to `useInferenceStatus`. No cache management UI
existed.
7. **Dead code accumulation** — Old `KokoroWorkerManager` (232 lines),
legacy Whisper message types, and scattered duplicate constants.

## Changes

### Phase 0 — Critical Performance & Bugs
- Add `fp16-webgpu` dtype for Kokoro TTS (~2x inference speed on
supported GPUs)
- Fix Whisper warm-up tensor from `[1, 128, 3000]` → `[1, 128, 1]`
(minimal shader compilation)
- Fix Whisper worker silent error bug (add try-catch to `generate()` and
`load()`)

### Phase 1 — Data Transfer & Caching
- Switch Kokoro audio to Float32Array transferable (skip WAV blob encode
in worker, lightweight WAV encode on main thread)
- Cache `listVoices` results (skip redundant model reload when adapter
state is `ready`)
- Normalize progress reporting to 0-100 across all adapters,
differentiate `warmup` phase

### Phase 2 — Protocol Unification & Infrastructure
- Migrate all 3 workers + 3 adapters to unified `protocol.ts` message
types (`load-model`, `run-inference`, `model-ready`, `inference-result`,
`progress`, `error`)
- Wire `GPUResourceCoordinator` into all adapters (VRAM allocation
tracking, LRU ordering, memory pressure events)
- Wire `LoadQueue` into all adapters (priority-based sequential model
loading: TTS=10 > ASR=5 > BG_REMOVAL=1)
- Add `coordinator.ts` global singleton for GPU coordinator + load queue
- Add WebGPU detection + WASM fallback in Whisper and Background Removal
workers

### Phase 3 — Error Recovery & Observability
- Add restart logic with exponential backoff to Whisper adapter
(matching Kokoro's existing pattern)
- Integrate `classifyError()` (OOM / DEVICE_LOST / TIMEOUT
classification) in Whisper adapter
- Extend `defaultPerfTracer` to Whisper `transcribe()` and Background
Removal `processImage()`
- Wire `useInferenceStatus` into all 3 adapters (downloading → ready →
terminated lifecycle)

### Phase 4 — Tests
- Add unit tests for `AsyncMutex` (4 tests), `LoadQueue` (4 tests),
`GPUResourceCoordinator` (7 tests) — all 15 passing

### Phase 5 — Cleanup & Features
- Delete old `KokoroWorkerManager` (232 lines, zero consumers)
- Delete orphaned `libs/workers/types.ts` (old Whisper message types)
- Clean up `workers/kokoro/types.ts` (remove legacy message types, keep
domain types)
- Create centralized `constants.ts` (MODEL_IDS, MODEL_NAMES, TIMEOUTS,
MAX_RESTARTS)
- Remove hardcoded WebGPU check from background-removal devtools pages
(worker auto-detects)
- Add `useModelPreload` composable for generic idle-time preloading
- Add `useInferencePreload` composable that reads provider config and
preloads configured local models
- Wire preloading into both `stage-web` and `stage-tamagotchi` App.vue
(Kokoro TTS preloads 3s after init)
- Add `ModelCacheManager.vue` settings component (cache size display,
per-model status, clear cache)
- Document GPU Device isolation architecture in protocol.ts

## After This PR

- All inference workers speak the same protocol → adding a new model
adapter is straightforward
- GPU memory is tracked across all models with automatic pressure
warnings at 80%/95% of VRAM budget
- Models load sequentially via priority queue → no bandwidth/VRAM
contention
- Workers auto-detect WebGPU and fall back to WASM → works on browsers
without WebGPU
- Kokoro TTS preloads during idle time → "instant" first use for
configured users
- All adapters auto-restart on worker crashes (max 3 attempts,
exponential backoff)
- 15 unit tests cover core infrastructure (mutex, queue, coordinator)
- Zero dead code remains in the inference pipeline

## Test Plan

- [x] `pnpm exec vitest run packages/stage-ui/src/libs/inference/` — 15
tests pass
- [x] `pnpm -F @proj-airi/stage-ui exec tsc --noEmit` — no TypeScript
errors
- [x] `pnpm lint:fix` — no lint errors in changed files
- [ ] Manual: verify Kokoro TTS works with fp16-webgpu on a supported
browser
- [ ] Manual: verify Whisper ASR loads and transcribes correctly
- [ ] Manual: verify Background Removal works in devtools page
- [ ] Manual: verify preloading triggers in console (`[Preload] Loading
kokoro-...`)

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-04-13 23:19:27 +08:00

210 lines
7.1 KiB
TypeScript

/// <reference types="./vite.config-env.d.ts" />
import type { PluginOption } from 'vite'
import process from 'node:process'
import { execSync } from 'node:child_process'
import { join, resolve } from 'node:path'
import VueI18n from '@intlify/unplugin-vue-i18n/vite'
import templateCompilerOptions from '@tresjs/core/template-compiler-options'
import Vue from '@vitejs/plugin-vue'
import Unocss from 'unocss/vite'
import Info from 'unplugin-info/vite'
import VueRouter from 'unplugin-vue-router/vite'
import Yaml from 'unplugin-yaml/vite'
import mkcert from 'vite-plugin-mkcert'
import VueDevTools from 'vite-plugin-vue-devtools'
import Layouts from 'vite-plugin-vue-layouts'
import VueMacros from 'vue-macros/vite'
import { tryCatch } from '@moeru/std'
import { Download } from '@proj-airi/unplugin-fetch/vite'
import { DownloadLive2DSDK } from '@proj-airi/unplugin-live2d-sdk/vite'
import { defineConfig } from 'vite'
// import { isEnvTruthy } from '@proj-airi/stage-shared'
function isEnvTruthy(value: string | undefined | null): boolean {
if (value == null)
return false
return /^(?:1|true|t|yes|y|on)$/i.test(value.trim())
}
const stageUIAssetsRoot = resolve(join(import.meta.dirname, '..', '..', 'packages', 'stage-ui', 'src', 'assets'))
const sharedCacheDir = resolve(join(import.meta.dirname, '..', '..', '.cache'))
export default defineConfig({
optimizeDeps: {
exclude: [
// Internal Packages
'@proj-airi/stage-ui/*',
'@proj-airi/drizzle-duckdb-wasm',
'@proj-airi/drizzle-duckdb-wasm/*',
// Static Assets: Models, Images, etc.
'public/assets/*',
// Live2D SDK
'@framework/live2dcubismframework',
'@framework/math/cubismmatrix44',
'@framework/type/csmvector',
'@framework/math/cubismviewmatrix',
'@framework/cubismdefaultparameterid',
'@framework/cubismmodelsettingjson',
'@framework/effect/cubismbreath',
'@framework/effect/cubismeyeblink',
'@framework/model/cubismusermodel',
'@framework/motion/acubismmotion',
'@framework/motion/cubismmotionqueuemanager',
'@framework/type/csmmap',
'@framework/utils/cubismdebug',
'@framework/model/cubismmoc',
],
},
resolve: {
alias: {
'@proj-airi/server-sdk': resolve(join(import.meta.dirname, '..', '..', 'packages', 'server-sdk', 'src')),
'@proj-airi/i18n': resolve(join(import.meta.dirname, '..', '..', 'packages', 'i18n', 'src')),
'@proj-airi/stage-ui': resolve(join(import.meta.dirname, '..', '..', 'packages', 'stage-ui', 'src')),
'@proj-airi/stage-layouts': resolve(join(import.meta.dirname, '..', '..', 'packages', 'stage-layouts', 'src')),
'@proj-airi/stage-pages': resolve(join(import.meta.dirname, '..', '..', 'packages', 'stage-pages', 'src')),
'@proj-airi/stage-shared': resolve(join(import.meta.dirname, '..', '..', 'packages', 'stage-shared', 'src')),
},
},
server: {
host: '0.0.0.0',
port: 5273,
fs: {
// To mute errors like:
// The request id ".../node_modules/@fontsource/sniglet/files/sniglet-latin-400-normal.woff" is outside of Vite serving allow list.
//
// See: https://vite.dev/config/server-options#server-fs-strict
strict: false,
},
warmup: {
clientFiles: [
`${resolve(join(import.meta.dirname, '..', '..', 'packages', 'stage-ui', 'src'))}/*.vue`,
`${resolve(join(import.meta.dirname, '..', '..', 'packages', 'stage-pages', 'src'))}/*.vue`,
`${resolve(join(import.meta.dirname, '..', '..', 'packages', 'stage-layouts', 'src'))}/*.vue`,
],
},
},
build: {
sourcemap: true,
},
worker: {
format: 'es',
rollupOptions: {
output: {
inlineDynamicImports: false,
},
},
},
plugins: [
...isEnvTruthy(process.env.VITE_SKIP_MKCERT ?? '')
? []
: [mkcert((() => {
// Workaround: plugin's bundled downloader has a feaxios bug, prefer system mkcert
const command = process.platform === 'win32' ? 'where' : 'which'
const { data } = tryCatch(() => ({ mkcertPath: execSync(`${command} mkcert`, { stdio: 'pipe' }).toString().trim().split(/\r?\n/)[0] }))
return data
})())],
Info(),
Yaml(),
VueMacros({
plugins: {
vue: Vue({
include: [/\.vue$/, /\.md$/],
...templateCompilerOptions,
}),
vueJsx: false,
},
betterDefine: false,
}),
// https://github.com/posva/unplugin-vue-router
VueRouter({
extensions: ['.vue', '.md'],
dts: resolve(import.meta.dirname, 'src/typed-router.d.ts'),
importMode: 'async',
routesFolder: [
resolve(import.meta.dirname, 'src', 'pages'),
{
src: resolve(import.meta.dirname, '..', '..', 'packages', 'stage-pages', 'src', 'pages'),
exclude: base => [
...base,
'**/settings/connection/index.vue',
],
},
],
exclude: ['**/components/**'],
}),
// https://github.com/JohnCampionJr/vite-plugin-vue-layouts
Layouts({
layoutsDirs: [
resolve(import.meta.dirname, 'src', 'layouts'),
resolve(import.meta.dirname, '..', '..', 'packages', 'stage-layouts', 'src', 'layouts'),
],
}),
// https://github.com/antfu/unocss
// see uno.config.ts for config
Unocss(),
// https://github.com/intlify/bundle-tools/tree/main/packages/unplugin-vue-i18n
VueI18n({
runtimeOnly: true,
compositionOnly: true,
fullInstall: true,
}),
// https://github.com/webfansplz/vite-plugin-vue-devtools
VueDevTools(),
DownloadLive2DSDK(),
Download('https://dist.ayaka.moe/live2d-models/hiyori_free_zh.zip', 'hiyori_free_zh.zip', 'live2d/models', { parentDir: stageUIAssetsRoot, cacheDir: sharedCacheDir }),
Download('https://dist.ayaka.moe/live2d-models/hiyori_pro_zh.zip', 'hiyori_pro_zh.zip', 'live2d/models', { parentDir: stageUIAssetsRoot, cacheDir: sharedCacheDir }),
Download('https://dist.ayaka.moe/vrm-models/VRoid-Hub/AvatarSample-A/AvatarSample_A.vrm', 'AvatarSample_A.vrm', 'vrm/models/AvatarSample-A', { parentDir: stageUIAssetsRoot, cacheDir: sharedCacheDir }),
Download('https://dist.ayaka.moe/vrm-models/VRoid-Hub/AvatarSample-B/AvatarSample_B.vrm', 'AvatarSample_B.vrm', 'vrm/models/AvatarSample-B', { parentDir: stageUIAssetsRoot, cacheDir: sharedCacheDir }),
...isEnvTruthy(process.env.VITE_CAP_SYNC_IOS_AFTER_BUILD ?? '')
? [{
name: 'proj-airi:capacitor-sync',
closeBundle: {
sequential: true,
handler() {
if (this.meta.watchMode) {
execSync('cap sync ios', { stdio: 'inherit' })
}
},
},
} as PluginOption]
: [],
{
name: 'proj-airi:defines',
config(ctx) {
const define: Record<string, any> = {
'import.meta.env.RUNTIME_ENVIRONMENT': '\'capacitor\'',
}
if (ctx.mode === 'development') {
define['import.meta.env.URL_MODE'] = '\'server\''
}
if (ctx.mode === 'production') {
define['import.meta.env.URL_MODE'] = '\'file\''
}
return { define }
},
},
],
})