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>
This commit is contained in:
NJX
co-authored by
Claude Opus 4.6
autofix-ci[bot]
Claude Opus 4.6
autofix-ci[bot]
parent
d86314f5a3
commit
147d077aa7
@@ -1,13 +1,10 @@
|
||||
<script setup lang="ts">
|
||||
import type { PreTrainedModel, Processor } from '@huggingface/transformers'
|
||||
import { createBackgroundRemovalAdapter } from '@proj-airi/stage-ui/libs/inference/adapters/background-removal'
|
||||
import { Button, Checkbox, InputFile } from '@proj-airi/ui'
|
||||
import { computed, onMounted, onUnmounted, ref, watch } from 'vue'
|
||||
|
||||
import { AutoModel, AutoProcessor, env, RawImage } from '@huggingface/transformers'
|
||||
import { Button, Checkbox, InputFileCard } from '@proj-airi/ui'
|
||||
import { check } from 'gpuu/webgpu'
|
||||
import { computed, onMounted, ref, watch } from 'vue'
|
||||
const adapter = createBackgroundRemovalAdapter()
|
||||
|
||||
const model = ref<PreTrainedModel>()
|
||||
const processor = ref<Processor>()
|
||||
const error = ref<unknown>()
|
||||
const loading = ref(true)
|
||||
const processing = ref(false)
|
||||
@@ -62,17 +59,8 @@ watch(autoProcess, (enabled) => {
|
||||
|
||||
onMounted(async () => {
|
||||
try {
|
||||
if (!((await check()).supported)) {
|
||||
throw new Error('WebGPU is not supported in this browser.')
|
||||
}
|
||||
|
||||
const model_id = 'Xenova/modnet'
|
||||
env.backends.onnx.wasm!.proxy = false
|
||||
model.value ??= await AutoModel.from_pretrained(model_id, {
|
||||
device: 'webgpu',
|
||||
})
|
||||
|
||||
processor.value ??= await AutoProcessor.from_pretrained(model_id, {})
|
||||
// Worker auto-detects WebGPU and falls back to WASM if unavailable
|
||||
await adapter.load()
|
||||
}
|
||||
catch (err) {
|
||||
error.value = err
|
||||
@@ -81,26 +69,26 @@ onMounted(async () => {
|
||||
loading.value = false
|
||||
})
|
||||
|
||||
onUnmounted(() => {
|
||||
adapter.terminate()
|
||||
})
|
||||
|
||||
async function processImage(item: ImageItem, index: number): Promise<void> {
|
||||
if (!model.value || !processor.value)
|
||||
if (adapter.state !== 'ready')
|
||||
return
|
||||
|
||||
try {
|
||||
item.status = 'processing'
|
||||
currentProcessingIndex.value = index
|
||||
|
||||
// Load image
|
||||
const img = await RawImage.fromURL(item.originalUrl)
|
||||
// Load image into a canvas to get ImageData
|
||||
const img = new Image()
|
||||
img.src = item.originalUrl
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
img.onload = () => resolve()
|
||||
img.onerror = () => reject(new Error('Failed to load image'))
|
||||
})
|
||||
|
||||
// Pre-process image
|
||||
const { pixel_values } = await processor.value(img)
|
||||
|
||||
// Predict alpha matte
|
||||
const { output } = await model.value({ input: pixel_values })
|
||||
|
||||
const maskData = (await RawImage.fromTensor(output[0].mul(255).to('uint8')).resize(img.width, img.height)).data
|
||||
|
||||
// Create new canvas
|
||||
const canvas = document.createElement('canvas')
|
||||
canvas.width = img.width
|
||||
canvas.height = img.height
|
||||
@@ -108,16 +96,14 @@ async function processImage(item: ImageItem, index: number): Promise<void> {
|
||||
if (!ctx)
|
||||
return
|
||||
|
||||
// Draw original image output to canvas
|
||||
ctx.drawImage(img.toCanvas(), 0, 0)
|
||||
ctx.drawImage(img, 0, 0)
|
||||
const imageData = ctx.getImageData(0, 0, img.width, img.height)
|
||||
|
||||
// Update alpha channel
|
||||
const pixelData = ctx.getImageData(0, 0, img.width, img.height)
|
||||
for (let j = 0; j < maskData.length; ++j) {
|
||||
pixelData.data[4 * j + 3] = maskData[j]
|
||||
}
|
||||
// Process in worker (off main thread!)
|
||||
const resultData = await adapter.processImage(imageData)
|
||||
|
||||
ctx.putImageData(pixelData, 0, 0)
|
||||
// Draw result to canvas for export
|
||||
ctx.putImageData(resultData, 0, 0)
|
||||
item.processedUrl = canvas.toDataURL('image/png')
|
||||
item.status = 'done'
|
||||
}
|
||||
@@ -127,7 +113,7 @@ async function processImage(item: ImageItem, index: number): Promise<void> {
|
||||
}
|
||||
|
||||
async function processAllImages() {
|
||||
if (!model.value || !processor.value || processing.value)
|
||||
if (adapter.state !== 'ready' || processing.value)
|
||||
return
|
||||
|
||||
processing.value = true
|
||||
@@ -222,7 +208,7 @@ function hidePreview() {
|
||||
<!-- Main content -->
|
||||
<template v-else>
|
||||
<!-- File upload area -->
|
||||
<InputFileCard v-model="imageFiles" accept="image/*" multiple w-full />
|
||||
<InputFile v-model="imageFiles" accept="image/*" multiple w-full />
|
||||
|
||||
<!-- Controls -->
|
||||
<div flex flex-wrap items-center justify-between gap-4>
|
||||
@@ -236,7 +222,7 @@ function hidePreview() {
|
||||
<Button
|
||||
v-if="pendingCount > 0"
|
||||
:label="processing ? `Processing... ${progressPercent}%` : `Process ${pendingCount} image${pendingCount > 1 ? 's' : ''}`"
|
||||
:disabled="processing || !model"
|
||||
:disabled="processing"
|
||||
:loading="processing"
|
||||
@click="processAllImages"
|
||||
/>
|
||||
|
||||
@@ -28,7 +28,7 @@ import { defineConfig } from 'vite'
|
||||
function isEnvTruthy(value: string | undefined | null): boolean {
|
||||
if (value == null)
|
||||
return false
|
||||
// eslint-disable-next-line e18e/prefer-static-regex
|
||||
|
||||
return /^(?:1|true|t|yes|y|on)$/i.test(value.trim())
|
||||
}
|
||||
|
||||
@@ -109,7 +109,7 @@ export default defineConfig({
|
||||
: [mkcert((() => {
|
||||
// Workaround: plugin's bundled downloader has a feaxios bug, prefer system mkcert
|
||||
const command = process.platform === 'win32' ? 'where' : 'which'
|
||||
// eslint-disable-next-line e18e/prefer-static-regex
|
||||
|
||||
const { data } = tryCatch(() => ({ mkcertPath: execSync(`${command} mkcert`, { stdio: 'pipe' }).toString().trim().split(/\r?\n/)[0] }))
|
||||
return data
|
||||
})())],
|
||||
|
||||
@@ -12,7 +12,7 @@ function hasXcode26OrAbove() {
|
||||
try {
|
||||
const output = execSync('xcodebuild -version')
|
||||
.toString()
|
||||
// eslint-disable-next-line e18e/prefer-static-regex
|
||||
|
||||
.match(/Xcode (\d+)/)
|
||||
if (!output)
|
||||
return false
|
||||
|
||||
@@ -25,7 +25,7 @@ function getContentType(pathname: string) {
|
||||
export async function startUpdateTestServer(options: { port: number, rootDir: string }) {
|
||||
const server = createServer(async (request, response) => {
|
||||
const pathname = request.url?.split('?')[0] || '/'
|
||||
// eslint-disable-next-line e18e/prefer-static-regex
|
||||
|
||||
const safePath = normalize(pathname).replace(/^(\.\.(\/|\\|$))+/, '')
|
||||
const filePath = join(options.rootDir, safePath === '/' ? '/index.html' : safePath)
|
||||
|
||||
|
||||
@@ -21,7 +21,6 @@ export async function getVersion(options: { release: boolean, autoTag: boolean,
|
||||
|
||||
// If --tag is specified, use the provided tag
|
||||
if (options.tag[0] !== 'true') {
|
||||
// eslint-disable-next-line e18e/prefer-static-regex
|
||||
version = String(options.tag[0]).replace(/^v/, '').trim()
|
||||
}
|
||||
// Otherwise, even for --tag option (true / enabled), ignore the input
|
||||
@@ -42,7 +41,7 @@ export async function getVersion(options: { release: boolean, autoTag: boolean,
|
||||
// fetch the latest git ref
|
||||
try {
|
||||
const res = await x('git', ['describe', '--tags', '--abbrev=0'])
|
||||
// eslint-disable-next-line e18e/prefer-static-regex
|
||||
|
||||
return String(res.stdout).replace(/^v/, '').trim()
|
||||
}
|
||||
catch {
|
||||
|
||||
@@ -3,6 +3,7 @@ import { defineInvokeHandler } from '@moeru/eventa'
|
||||
import { useElectronEventaContext, useElectronEventaInvoke } from '@proj-airi/electron-vueuse'
|
||||
import { themeColorFromValue, useThemeColor } from '@proj-airi/stage-layouts/composables/theme-color'
|
||||
import { ToasterRoot } from '@proj-airi/stage-ui/components'
|
||||
import { useInferencePreload } from '@proj-airi/stage-ui/composables'
|
||||
import { useSharedAnalyticsStore } from '@proj-airi/stage-ui/stores/analytics'
|
||||
import { useCharacterOrchestratorStore } from '@proj-airi/stage-ui/stores/character'
|
||||
import { useChatSessionStore } from '@proj-airi/stage-ui/stores/chat/session-store'
|
||||
@@ -60,6 +61,7 @@ const chatSessionStore = useChatSessionStore()
|
||||
const serverChannelStore = useModsServerChannelStore()
|
||||
const characterOrchestratorStore = useCharacterOrchestratorStore()
|
||||
const analyticsStore = useSharedAnalyticsStore()
|
||||
const inferencePreload = useInferencePreload()
|
||||
const pluginHostInspectorStore = usePluginHostInspectorStore()
|
||||
const stageWindowLifecycleStore = useStageWindowLifecycleStore()
|
||||
const settingsAudioDeviceStore = useSettingsAudioDevice()
|
||||
@@ -157,6 +159,9 @@ onMounted(async () => {
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// Preload local inference models (Kokoro TTS, etc.) in background after a delay
|
||||
inferencePreload.triggerPreload()
|
||||
})
|
||||
|
||||
watch(themeColorsHue, () => {
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
<script setup lang="ts">
|
||||
import { OnboardingDialog, OnboardingStepAnalyticsNotice, ToasterRoot } from '@proj-airi/stage-ui/components'
|
||||
import { useInferencePreload } from '@proj-airi/stage-ui/composables'
|
||||
import { isPosthogAvailableInBuild, useSharedAnalyticsStore } from '@proj-airi/stage-ui/stores/analytics'
|
||||
import { useCharacterOrchestratorStore } from '@proj-airi/stage-ui/stores/character'
|
||||
import { useChatSessionStore } from '@proj-airi/stage-ui/stores/chat/session-store'
|
||||
@@ -37,6 +38,7 @@ const { showingSetup } = storeToRefs(onboardingStore)
|
||||
const { isDark } = useTheme()
|
||||
const cardStore = useAiriCardStore()
|
||||
const analyticsStore = useSharedAnalyticsStore()
|
||||
const inferencePreload = useInferencePreload()
|
||||
|
||||
const primaryColor = computed(() => {
|
||||
return isDark.value
|
||||
@@ -96,6 +98,9 @@ onMounted(async () => {
|
||||
await displayModelsStore.loadDisplayModelsFromIndexedDB()
|
||||
await settingsStore.initializeStageModel()
|
||||
await settingsAudioDeviceStore.initialize()
|
||||
|
||||
// Preload local inference models (Kokoro TTS, etc.) in background after a delay
|
||||
inferencePreload.triggerPreload()
|
||||
})
|
||||
|
||||
onUnmounted(() => {
|
||||
|
||||
@@ -1,13 +1,10 @@
|
||||
<script setup lang="ts">
|
||||
import type { PreTrainedModel, Processor } from '@huggingface/transformers'
|
||||
import { createBackgroundRemovalAdapter } from '@proj-airi/stage-ui/libs/inference/adapters/background-removal'
|
||||
import { Button, Checkbox, InputFile } from '@proj-airi/ui'
|
||||
import { computed, onMounted, onUnmounted, ref, watch } from 'vue'
|
||||
|
||||
import { AutoModel, AutoProcessor, env, RawImage } from '@huggingface/transformers'
|
||||
import { Button, Checkbox, InputFileCard } from '@proj-airi/ui'
|
||||
import { check } from 'gpuu/webgpu'
|
||||
import { computed, onMounted, ref, watch } from 'vue'
|
||||
const adapter = createBackgroundRemovalAdapter()
|
||||
|
||||
const model = ref<PreTrainedModel>()
|
||||
const processor = ref<Processor>()
|
||||
const error = ref<unknown>()
|
||||
const loading = ref(true)
|
||||
const processing = ref(false)
|
||||
@@ -62,17 +59,8 @@ watch(autoProcess, (enabled) => {
|
||||
|
||||
onMounted(async () => {
|
||||
try {
|
||||
if (!((await check()).supported)) {
|
||||
throw new Error('WebGPU is not supported in this browser.')
|
||||
}
|
||||
|
||||
const model_id = 'Xenova/modnet'
|
||||
env.backends.onnx.wasm!.proxy = false
|
||||
model.value ??= await AutoModel.from_pretrained(model_id, {
|
||||
device: 'webgpu',
|
||||
})
|
||||
|
||||
processor.value ??= await AutoProcessor.from_pretrained(model_id, {})
|
||||
// Worker auto-detects WebGPU and falls back to WASM if unavailable
|
||||
await adapter.load()
|
||||
}
|
||||
catch (err) {
|
||||
error.value = err
|
||||
@@ -81,26 +69,26 @@ onMounted(async () => {
|
||||
loading.value = false
|
||||
})
|
||||
|
||||
onUnmounted(() => {
|
||||
adapter.terminate()
|
||||
})
|
||||
|
||||
async function processImage(item: ImageItem, index: number): Promise<void> {
|
||||
if (!model.value || !processor.value)
|
||||
if (adapter.state !== 'ready')
|
||||
return
|
||||
|
||||
try {
|
||||
item.status = 'processing'
|
||||
currentProcessingIndex.value = index
|
||||
|
||||
// Load image
|
||||
const img = await RawImage.fromURL(item.originalUrl)
|
||||
// Load image into a canvas to get ImageData
|
||||
const img = new Image()
|
||||
img.src = item.originalUrl
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
img.onload = () => resolve()
|
||||
img.onerror = () => reject(new Error('Failed to load image'))
|
||||
})
|
||||
|
||||
// Pre-process image
|
||||
const { pixel_values } = await processor.value(img)
|
||||
|
||||
// Predict alpha matte
|
||||
const { output } = await model.value({ input: pixel_values })
|
||||
|
||||
const maskData = (await RawImage.fromTensor(output[0].mul(255).to('uint8')).resize(img.width, img.height)).data
|
||||
|
||||
// Create new canvas
|
||||
const canvas = document.createElement('canvas')
|
||||
canvas.width = img.width
|
||||
canvas.height = img.height
|
||||
@@ -108,16 +96,14 @@ async function processImage(item: ImageItem, index: number): Promise<void> {
|
||||
if (!ctx)
|
||||
return
|
||||
|
||||
// Draw original image output to canvas
|
||||
ctx.drawImage(img.toCanvas(), 0, 0)
|
||||
ctx.drawImage(img, 0, 0)
|
||||
const imageData = ctx.getImageData(0, 0, img.width, img.height)
|
||||
|
||||
// Update alpha channel
|
||||
const pixelData = ctx.getImageData(0, 0, img.width, img.height)
|
||||
for (let j = 0; j < maskData.length; ++j) {
|
||||
pixelData.data[4 * j + 3] = maskData[j]
|
||||
}
|
||||
// Process in worker (off main thread!)
|
||||
const resultData = await adapter.processImage(imageData)
|
||||
|
||||
ctx.putImageData(pixelData, 0, 0)
|
||||
// Draw result to canvas for export
|
||||
ctx.putImageData(resultData, 0, 0)
|
||||
item.processedUrl = canvas.toDataURL('image/png')
|
||||
item.status = 'done'
|
||||
}
|
||||
@@ -127,7 +113,7 @@ async function processImage(item: ImageItem, index: number): Promise<void> {
|
||||
}
|
||||
|
||||
async function processAllImages() {
|
||||
if (!model.value || !processor.value || processing.value)
|
||||
if (adapter.state !== 'ready' || processing.value)
|
||||
return
|
||||
|
||||
processing.value = true
|
||||
@@ -222,7 +208,7 @@ function hidePreview() {
|
||||
<!-- Main content -->
|
||||
<template v-else>
|
||||
<!-- File upload area -->
|
||||
<InputFileCard v-model="imageFiles" accept="image/*" multiple w-full />
|
||||
<InputFile v-model="imageFiles" accept="image/*" multiple w-full />
|
||||
|
||||
<!-- Controls -->
|
||||
<div flex flex-wrap items-center justify-between gap-4>
|
||||
@@ -236,7 +222,7 @@ function hidePreview() {
|
||||
<Button
|
||||
v-if="pendingCount > 0"
|
||||
:label="processing ? `Processing... ${progressPercent}%` : `Process ${pendingCount} image${pendingCount > 1 ? 's' : ''}`"
|
||||
:disabled="processing || !model"
|
||||
:disabled="processing"
|
||||
:loading="processing"
|
||||
@click="processAllImages"
|
||||
/>
|
||||
|
||||
@@ -108,7 +108,7 @@ export default defineConfig({
|
||||
? [Mkcert((() => {
|
||||
// Workaround: plugin's bundled downloader has a feaxios bug, prefer system mkcert
|
||||
const command = process.platform === 'win32' ? 'where' : 'which'
|
||||
// eslint-disable-next-line e18e/prefer-static-regex
|
||||
|
||||
const { data } = tryCatch(() => ({ mkcertPath: execSync(`${command} mkcert`, { stdio: 'pipe' }).toString().trim().split(/\r?\n/)[0] }))
|
||||
return data
|
||||
})())]
|
||||
|
||||
Reference in New Issue
Block a user