concept-collection / voicenote
voicenote / models.js
90 lines · 3.0 KBBlameHistoryRaw
1// Speech models offered in the picker, and the per-device weights they need.
2//
3// The two backends want different quantizations, so they download different
4// files and the picker has to quote the size for whichever one will actually
5// run:
6//
7// wasm — int8 encoder + int8 merged decoder. Small, CPU-friendly.
8// webgpu — fp32 encoder + 4-bit merged decoder. This is the combination the
9// transformers.js WebGPU examples use; int8 matmul is poorly served
10// on WebGPU. It runs far faster but downloads 2–3× as much, because
11// the 4-bit decoder only quantizes matmul weights and leaves the
12// embeddings at full precision.
13//
14// `mb` figures are the encoder + decoder + tokenizer/config bytes, taken from
15// the HuggingFace file listing. Everything is cached after the first visit.
16export const MODELS = [
17 {
18 id: 'onnx-community/moonshine-tiny-ONNX',
19 label: 'Moonshine Tiny',
20 mb: { wasm: 32, webgpu: 79 },
21 note: 'fastest; lighter punctuation',
22 },
23 {
24 id: 'onnx-community/whisper-tiny.en',
25 label: 'Whisper Tiny',
26 mb: { wasm: 44, webgpu: 122 },
27 note: 'quick; less accurate on hard words',
28 },
29 {
30 id: 'onnx-community/moonshine-base-ONNX',
31 label: 'Moonshine Base',
32 mb: { wasm: 67, webgpu: 157 },
33 note: 'fast and accurate; lighter punctuation',
34 },
35 {
36 id: 'onnx-community/whisper-base.en',
37 label: 'Whisper Base',
38 mb: { wasm: 80, webgpu: 209 },
39 note: 'good accuracy and punctuation',
40 },
41 {
42 id: 'onnx-community/whisper-small.en',
43 label: 'Whisper Small',
44 mb: { wasm: 252, webgpu: 588 },
45 note: 'most accurate; slow without WebGPU',
46 },
47];
49export const DEFAULT_MODEL = 'onnx-community/whisper-base.en';
51// Quantization suffixes per device — these decide both which files
52// transformers.js fetches and which ones we look for in the cache.
53export const SUFFIX = {
54 wasm: { encoder: '_quantized', decoder: '_quantized' },
55 webgpu: { encoder: '', decoder: '_q4' },
56};
58export async function detectDevice() {
59 if (!navigator.gpu) return 'wasm';
60 try {
61 return (await navigator.gpu.requestAdapter()) ? 'webgpu' : 'wasm';
62 } catch {
63 return 'wasm';
64 }
67// The two large files for a model on a given device. Their presence in the
68// cache is what "downloaded" in the picker means.
69const weightFiles = (id, device) => {
70 const s = SUFFIX[device];
71 return [
72 `https://huggingface.co/${id}/resolve/main/onnx/encoder_model${s.encoder}.onnx`,
73 `https://huggingface.co/${id}/resolve/main/onnx/decoder_model_merged${s.decoder}.onnx`,
74 ];
75};
77export async function cachedModels(device) {
78 const cached = new Set();
79 if (!self.caches) return cached;
80 try {
81 const cache = await caches.open('transformers-cache');
82 await Promise.all(
83 MODELS.map(async (m) => {
84 const hits = await Promise.all(weightFiles(m.id, device).map((u) => cache.match(u)));
85 if (hits.every(Boolean)) cached.add(m.id);
86 }),
87 );
88 } catch { /* cache unavailable (private mode, etc.) — report none cached */ }
89 return cached;