Add download progress tracking

decoder_worker.js

@@ -0,0 +1,353 @@
+/**
+ * Decoder Worker - Runs adapter + decoder in a separate thread
+ */
+
+importScripts('https://cdn.jsdelivr.net/npm/onnxruntime-web@1.17.0/dist/ort.min.js');
+
+// Configure ONNX Runtime to find WASM files from CDN
+ort.env.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.17.0/dist/';
+
+// Helper to fetch model with progress reporting
+async function fetchModelWithProgress(url, modelName) {
+    const response = await fetch(url);
+    if (!response.ok) {
+        throw new Error(`Failed to fetch ${modelName}: ${response.status}`);
+    }
+
+    const contentLength = response.headers.get('Content-Length');
+    const total = contentLength ? parseInt(contentLength, 10) : 0;
+
+    if (!response.body || !total) {
+        // No streaming support or unknown size - just download
+        const buffer = await response.arrayBuffer();
+        self.postMessage({
+            type: 'progress',
+            model: modelName,
+            loaded: buffer.byteLength,
+            total: buffer.byteLength,
+            done: true
+        });
+        return buffer;
+    }
+
+    const reader = response.body.getReader();
+    const chunks = [];
+    let loaded = 0;
+
+    while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+
+        chunks.push(value);
+        loaded += value.length;
+
+        self.postMessage({
+            type: 'progress',
+            model: modelName,
+            loaded,
+            total,
+            done: false
+        });
+    }
+
+    self.postMessage({
+        type: 'progress',
+        model: modelName,
+        loaded: total,
+        total,
+        done: true
+    });
+
+    // Combine chunks into single ArrayBuffer
+    const result = new Uint8Array(loaded);
+    let offset = 0;
+    for (const chunk of chunks) {
+        result.set(chunk, offset);
+        offset += chunk.length;
+    }
+
+    return result.buffer;
+}
+
+// Model config
+let cfg = null;
+let tailLatency = 0;
+
+// Sessions
+let adapterSession = null;
+let decoderInitSession = null;
+let decoderStepSession = null;
+
+// Decoder state
+let crossCache = null;
+let selfCache = null;
+
+// Tokenizer
+let tokenizer = null;
+
+// Accumulated features
+let accumulatedFeatures = null;
+let currentSegmentId = null;
+
+class MoonshineTokenizer {
+    constructor() {
+        this.decoder = null;
+        this.vocab = null;
+    }
+
+    load(tokenizerJson) {
+        this.vocab = tokenizerJson.model.vocab;
+        this.decoder = Object.fromEntries(
+            Object.entries(this.vocab).map(([k, v]) => [v, k])
+        );
+    }
+
+    decode(tokenIds, skipSpecial = true) {
+        const specialTokens = new Set([0, 1, 2]);
+        let text = '';
+
+        for (const id of tokenIds) {
+            if (skipSpecial && specialTokens.has(id)) continue;
+            const token = this.decoder[id] || '';
+            text += token;
+        }
+
+        // Handle various space placeholder representations
+        text = text.replace(/\u0120/g, ' ');  // Ġ (GPT-2 style)
+        text = text.replace(/Ġ/g, ' ');       // Literal Ġ character
+        text = text.replace(/▁/g, ' ');       // SentencePiece style (U+2581)
+        text = text.replace(/\u010a/g, '\n'); // Newline marker
+
+        return text.trim();
+    }
+}
+
+async function runAdapter(features, dims) {
+    const feeds = {
+        'encoder_output': new ort.Tensor('float32', features, dims)
+    };
+    const results = await adapterSession.run(feeds);
+    return results.context;
+}
+
+async function initDecoderCache(context) {
+    const feeds = { 'context': context };
+    const results = await decoderInitSession.run(feeds);
+
+    // Store cross-attention cache (even-indexed layers)
+    crossCache = [];
+    for (let i = 0; i < cfg.depth * 2; i++) {
+        if ((i + 1) % 2 === 0) {
+            crossCache.push({
+                k: results[`cache_${i}_k`],
+                v: results[`cache_${i}_v`]
+            });
+        }
+    }
+
+    // Initialize empty self-attention cache
+    selfCache = [];
+    for (let i = 0; i < cfg.depth; i++) {
+        selfCache.push({
+            k: new ort.Tensor('float32', new Float32Array(0), [1, cfg.nheads, 0, cfg.head_dim]),
+            v: new ort.Tensor('float32', new Float32Array(0), [1, cfg.nheads, 0, cfg.head_dim])
+        });
+    }
+}
+
+async function decodeStep(tokenId, position) {
+    const feeds = {
+        'token_id': new ort.Tensor('int64', BigInt64Array.from([BigInt(tokenId)]), [1, 1]),
+        'position': new ort.Tensor('int64', BigInt64Array.from([BigInt(position)]), [1])
+    };
+
+    // Add cache inputs
+    let selfIdx = 0;
+    let crossIdx = 0;
+    for (let i = 0; i < cfg.depth * 2; i++) {
+        if ((i + 1) % 2 !== 0) {
+            feeds[`in_cache_${i}_k`] = selfCache[selfIdx].k;
+            feeds[`in_cache_${i}_v`] = selfCache[selfIdx].v;
+            selfIdx++;
+        } else {
+            feeds[`in_cache_${i}_k`] = crossCache[crossIdx].k;
+            feeds[`in_cache_${i}_v`] = crossCache[crossIdx].v;
+            crossIdx++;
+        }
+    }
+
+    const results = await decoderStepSession.run(feeds);
+
+    // Update self-attention cache
+    selfIdx = 0;
+    for (let i = 0; i < cfg.depth * 2; i++) {
+        if ((i + 1) % 2 !== 0) {
+            selfCache[selfIdx] = {
+                k: results[`out_cache_${i}_k`],
+                v: results[`out_cache_${i}_v`]
+            };
+            selfIdx++;
+        }
+    }
+
+    return results.logits;
+}
+
+async function decodeAccumulated() {
+    if (!accumulatedFeatures || accumulatedFeatures.dims[1] === 0) {
+        return '';
+    }
+
+    try {
+        const context = await runAdapter(accumulatedFeatures.data, accumulatedFeatures.dims);
+        await initDecoderCache(context);
+
+        const numFrames = accumulatedFeatures.dims[1];
+        const maxTokens = Math.max(10, Math.floor(numFrames * 1.5));
+
+        const tokens = [1];  // BOS
+        for (let step = 0; step < maxTokens; step++) {
+            const logits = await decodeStep(tokens[tokens.length - 1], step);
+
+            let maxIdx = 0;
+            let maxVal = logits.data[0];
+            for (let i = 1; i < cfg.vocab_size; i++) {
+                if (logits.data[i] > maxVal) {
+                    maxVal = logits.data[i];
+                    maxIdx = i;
+                }
+            }
+
+            tokens.push(maxIdx);
+            if (maxIdx === 2) break;  // EOS
+        }
+
+        return tokenizer.decode(tokens, true);
+    } catch (e) {
+        console.error('Decode error:', e);
+        return '';
+    }
+}
+
+self.onmessage = async function(e) {
+    const { type, data } = e.data;
+
+    switch (type) {
+        case 'init': {
+            try {
+                cfg = data.cfg;
+                const onnxUrl = data.onnxUrl;
+                const modelName = data.modelName;
+                const dtype = 'fp32';
+
+                tailLatency = cfg.n_future * cfg.encoder_depth;
+
+                // Load tokenizer
+                self.postMessage({ type: 'status', message: 'Loading tokenizer...' });
+                self.postMessage({ type: 'model_start', model: 'Tokenizer' });
+                const tokenizerResponse = await fetch(`${onnxUrl}/tokenizer.json`);
+                const tokenizerJson = await tokenizerResponse.json();
+                tokenizer = new MoonshineTokenizer();
+                tokenizer.load(tokenizerJson);
+                self.postMessage({ type: 'model_done', model: 'Tokenizer' });
+
+                // Initialize adapter
+                const adapterUrl = `${onnxUrl}/adapter_${modelName}_${dtype}.onnx`;
+                self.postMessage({ type: 'status', message: 'Loading adapter...' });
+                self.postMessage({ type: 'model_start', model: 'Adapter' });
+                const adapterBuffer = await fetchModelWithProgress(adapterUrl, 'Adapter');
+                adapterSession = await ort.InferenceSession.create(adapterBuffer);
+                self.postMessage({ type: 'model_done', model: 'Adapter' });
+
+                // Initialize decoder init
+                const decInitUrl = `${onnxUrl}/decoder_init_${modelName}_${dtype}.onnx`;
+                self.postMessage({ type: 'status', message: 'Loading decoder (init)...' });
+                self.postMessage({ type: 'model_start', model: 'Decoder Init' });
+                const decInitBuffer = await fetchModelWithProgress(decInitUrl, 'Decoder Init');
+                decoderInitSession = await ort.InferenceSession.create(decInitBuffer);
+                self.postMessage({ type: 'model_done', model: 'Decoder Init' });
+
+                // Initialize decoder step
+                const decStepUrl = `${onnxUrl}/decoder_step_${modelName}_${dtype}.onnx`;
+                self.postMessage({ type: 'status', message: 'Loading decoder (step)...' });
+                self.postMessage({ type: 'model_start', model: 'Decoder Step' });
+                const decStepBuffer = await fetchModelWithProgress(decStepUrl, 'Decoder Step');
+                decoderStepSession = await ort.InferenceSession.create(decStepBuffer);
+                self.postMessage({ type: 'model_done', model: 'Decoder Step' });
+
+                self.postMessage({ type: 'ready' });
+            } catch (err) {
+                self.postMessage({ type: 'error', message: err.message });
+            }
+            break;
+        }
+
+        case 'segment_start': {
+            accumulatedFeatures = null;
+            currentSegmentId = data.segmentId;
+            self.postMessage({ type: 'live_caption', text: '' });
+            break;
+        }
+
+        case 'segment_end': {
+            if (data.segmentId !== currentSegmentId) break;
+
+            const text = await decodeAccumulated();
+            self.postMessage({
+                type: 'transcript',
+                segmentId: data.segmentId,
+                text: text
+            });
+
+            accumulatedFeatures = null;
+            currentSegmentId = null;
+            self.postMessage({ type: 'live_caption', text: '' });
+            break;
+        }
+
+        case 'features': {
+            if (data.segmentId !== currentSegmentId) break;
+
+            const newFeatures = {
+                data: new Float32Array(data.features),
+                dims: data.dims
+            };
+
+            console.log(`Decoder received ${data.dims[1]} frames, accumulated: ${accumulatedFeatures ? accumulatedFeatures.dims[1] : 0}`);
+
+            if (accumulatedFeatures === null) {
+                accumulatedFeatures = newFeatures;
+            } else {
+                // Trim last tailLatency frames from accumulated
+                const numFrames = accumulatedFeatures.dims[1];
+                const keepFrames = Math.max(0, numFrames - tailLatency);
+
+                if (keepFrames > 0) {
+                    const totalFrames = keepFrames + newFeatures.dims[1];
+                    const combined = new Float32Array(totalFrames * cfg.dim);
+
+                    // Copy kept frames
+                    for (let f = 0; f < keepFrames; f++) {
+                        for (let d = 0; d < cfg.dim; d++) {
+                            combined[f * cfg.dim + d] = accumulatedFeatures.data[f * cfg.dim + d];
+                        }
+                    }
+                    // Copy new frames
+                    combined.set(newFeatures.data, keepFrames * cfg.dim);
+
+                    accumulatedFeatures = {
+                        data: combined,
+                        dims: [1, totalFrames, cfg.dim]
+                    };
+                } else {
+                    accumulatedFeatures = newFeatures;
+                }
+            }
+
+            // Live caption
+            const partialText = await decodeAccumulated();
+            self.postMessage({ type: 'live_caption', text: partialText });
+            break;
+        }
+    }
+};

encoder_worker.js

@@ -7,6 +7,68 @@ importScripts('https://cdn.jsdelivr.net/npm/onnxruntime-web@1.17.0/dist/ort.min.
 // Configure ONNX Runtime to find WASM files from CDN
 ort.env.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.17.0/dist/';
 
+// Helper to fetch model with progress reporting
+async function fetchModelWithProgress(url, modelName) {
+    const response = await fetch(url);
+    if (!response.ok) {
+        throw new Error(`Failed to fetch ${modelName}: ${response.status}`);
+    }
+
+    const contentLength = response.headers.get('Content-Length');
+    const total = contentLength ? parseInt(contentLength, 10) : 0;
+
+    if (!response.body || !total) {
+        // No streaming support or unknown size - just download
+        const buffer = await response.arrayBuffer();
+        self.postMessage({
+            type: 'progress',
+            model: modelName,
+            loaded: buffer.byteLength,
+            total: buffer.byteLength,
+            done: true
+        });
+        return buffer;
+    }
+
+    const reader = response.body.getReader();
+    const chunks = [];
+    let loaded = 0;
+
+    while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+
+        chunks.push(value);
+        loaded += value.length;
+
+        self.postMessage({
+            type: 'progress',
+            model: modelName,
+            loaded,
+            total,
+            done: false
+        });
+    }
+
+    self.postMessage({
+        type: 'progress',
+        model: modelName,
+        loaded: total,
+        total,
+        done: true
+    });
+
+    // Combine chunks into single ArrayBuffer
+    const result = new Uint8Array(loaded);
+    let offset = 0;
+    for (const chunk of chunks) {
+        result.set(chunk, offset);
+        offset += chunk.length;
+    }
+
+    return result.buffer;
+}
+
 // Model config
 let cfg = null;
 let preprocessor = null;
@@ -139,20 +201,26 @@ self.onmessage = async function(e) {
                 tailLatency = cfg.n_future * cfg.encoder_depth;
 
                 // Initialize preprocessor
+                const prepUrl = `${onnxUrl}/preprocessor_streaming_${modelName}_${dtype}.onnx`;
                 self.postMessage({ type: 'status', message: 'Loading preprocessor...' });
-                prepSession = await ort.InferenceSession.create(
-                    `${onnxUrl}/preprocessor_streaming_${modelName}_${dtype}.onnx`
-                );
+                self.postMessage({ type: 'model_start', model: 'Preprocessor' });
+                const prepBuffer = await fetchModelWithProgress(prepUrl, 'Preprocessor');
+                prepSession = await ort.InferenceSession.create(prepBuffer);
+                self.postMessage({ type: 'model_done', model: 'Preprocessor' });
+
                 prepDim = cfg.dim;
                 prepC1 = 2 * cfg.dim;
                 prepStateC1 = new Float32Array(4 * cfg.dim);
                 prepStateC2 = new Float32Array(4 * prepC1);
 
                 // Initialize encoder
+                const encUrl = `${onnxUrl}/encoder_${modelName}_${dtype}.onnx`;
                 self.postMessage({ type: 'status', message: 'Loading encoder...' });
-                encSession = await ort.InferenceSession.create(
-                    `${onnxUrl}/encoder_${modelName}_${dtype}.onnx`
-                );
+                self.postMessage({ type: 'model_start', model: 'Encoder' });
+                const encBuffer = await fetchModelWithProgress(encUrl, 'Encoder');
+                encSession = await ort.InferenceSession.create(encBuffer);
+                self.postMessage({ type: 'model_done', model: 'Encoder' });
+
                 encDim = cfg.dim;
                 encNPast = cfg.n_past;
                 encNFuture = cfg.n_future;

index.html

@@ -325,6 +325,43 @@
 
         .loading-text {
             color: #00d4ff;
+            font-size: 18px;
+            margin-bottom: 20px;
+        }
+
+        .loading-progress {
+            width: 300px;
+            margin: 0 auto;
+        }
+
+        .loading-progress-bar {
+            height: 8px;
+            background: #333;
+            border-radius: 4px;
+            overflow: hidden;
+            margin-bottom: 10px;
+        }
+
+        .loading-progress-fill {
+            height: 100%;
+            background: linear-gradient(90deg, #00d4ff, #00ff88);
+            width: 0%;
+            transition: width 0.3s ease;
+            border-radius: 4px;
+        }
+
+        .loading-progress-text {
+            font-size: 13px;
+            color: #888;
+            margin-bottom: 15px;
+        }
+
+        .loading-details {
+            font-size: 12px;
+            color: #666;
+            font-family: monospace;
+            max-height: 60px;
+            overflow: hidden;
         }
 
         .error-message {
@@ -346,6 +383,13 @@
         <div class="loading-content">
             <div class="loading-spinner"></div>
             <div class="loading-text" id="loadingText">Loading models...</div>
+            <div class="loading-progress">
+                <div class="loading-progress-bar">
+                    <div class="loading-progress-fill" id="loadingProgressFill"></div>
+                </div>
+                <div class="loading-progress-text" id="loadingProgressText">0 / 7 models</div>
+            </div>
+            <div class="loading-details" id="loadingDetails"></div>
         </div>
     </div>
 

streaming_asr.js

@@ -231,36 +231,73 @@ class PipelinedStreamingASR {
         this.onQueueUpdate = null;
     }
 
-    async loadModels(progressCallback) {
+    async loadModels(progressCallback, detailedProgressCallback) {
+        // Track overall progress
+        const totalModels = 7;  // VAD, Preprocessor, Encoder, Tokenizer, Adapter, Decoder Init, Decoder Step
+        let completedModels = 0;
+        let currentModel = '';
+        let currentProgress = { loaded: 0, total: 0 };
+
+        const updateProgress = () => {
+            const overallPercent = (completedModels / totalModels) * 100;
+            detailedProgressCallback?.({
+                completedModels,
+                totalModels,
+                overallPercent,
+                currentModel,
+                currentProgress
+            });
+        };
+
         // Initialize VAD
         try {
+            currentModel = 'VAD';
             progressCallback?.('Loading TenVAD...');
+            updateProgress();
             this.vad = new TenVAD(VAD_CHUNK_SAMPLES, 0.5);
             await this.vad.init('./ten_vad.js');
             console.log('Using TenVAD');
+            completedModels++;
+            updateProgress();
         } catch (e) {
             console.warn('TenVAD failed, using SimpleVAD:', e.message);
             this.vad = new SimpleVAD(SAMPLE_RATE, VAD_CHUNK_SAMPLES);
             await this.vad.init();
+            completedModels++;
+            updateProgress();
         }
 
-        // Initialize Encoder Worker
-        progressCallback?.('Loading encoder...');
-        await this.initEncoderWorker();
+        // Initialize Encoder Worker with progress tracking
+        progressCallback?.('Loading encoder models...');
+        await this.initEncoderWorker((model, progress) => {
+            currentModel = model;
+            currentProgress = progress;
+            updateProgress();
+        }, () => {
+            completedModels++;
+            updateProgress();
+        });
 
-        // Initialize Decoder Worker
-        progressCallback?.('Loading decoder...');
-        await this.initDecoderWorker();
+        // Initialize Decoder Worker with progress tracking
+        progressCallback?.('Loading decoder models...');
+        await this.initDecoderWorker((model, progress) => {
+            currentModel = model;
+            currentProgress = progress;
+            updateProgress();
+        }, () => {
+            completedModels++;
+            updateProgress();
+        });
 
         progressCallback?.('Ready!');
     }
 
-    initEncoderWorker() {
+    initEncoderWorker(onProgress, onModelDone) {
         return new Promise((resolve, reject) => {
             this.encoderWorker = new Worker('./encoder_worker.js');
 
             this.encoderWorker.onmessage = (e) => {
-                const { type, data } = e.data;
+                const { type } = e.data;
 
                 switch (type) {
                     case 'ready':
@@ -273,6 +310,12 @@ class PipelinedStreamingASR {
                     case 'status':
                         // Progress update from worker
                         break;
+                    case 'progress':
+                        onProgress?.(e.data.model, { loaded: e.data.loaded, total: e.data.total });
+                        break;
+                    case 'model_done':
+                        onModelDone?.(e.data.model);
+                        break;
                     case 'segment_start':
                         this.decoderWorker?.postMessage({ type: 'segment_start', data: { segmentId: e.data.segmentId } });
                         break;
@@ -304,7 +347,7 @@ class PipelinedStreamingASR {
         });
     }
 
-    initDecoderWorker() {
+    initDecoderWorker(onProgress, onModelDone) {
         return new Promise((resolve, reject) => {
             this.decoderWorker = new Worker('./decoder_worker.js');
 
@@ -321,6 +364,12 @@ class PipelinedStreamingASR {
                         break;
                     case 'status':
                         break;
+                    case 'progress':
+                        onProgress?.(e.data.model, { loaded: e.data.loaded, total: e.data.total });
+                        break;
+                    case 'model_done':
+                        onModelDone?.(e.data.model);
+                        break;
                     case 'transcript':
                         this.onTranscript?.(e.data.text, e.data.segmentId);
                         break;
@@ -610,6 +659,9 @@ class ASRDemoUI {
     initElements() {
         this.loadingOverlay = document.getElementById('loadingOverlay');
         this.loadingText = document.getElementById('loadingText');
+        this.loadingProgressFill = document.getElementById('loadingProgressFill');
+        this.loadingProgressText = document.getElementById('loadingProgressText');
+        this.loadingDetails = document.getElementById('loadingDetails');
         this.errorMessage = document.getElementById('errorMessage');
         this.statusDot = document.getElementById('statusDot');
         this.statusText = document.getElementById('statusText');
@@ -662,9 +714,14 @@ class ASRDemoUI {
             this.asr.onLiveCaption = (text) => this.updateLiveCaption(text);
             this.asr.onStatusUpdate = (status, text) => this.updateStatus(status, text);
 
-            await this.asr.loadModels((text) => {
-                this.loadingText.textContent = text;
-            });
+            await this.asr.loadModels(
+                (text) => {
+                    this.loadingText.textContent = text;
+                },
+                (progress) => {
+                    this.updateLoadingProgress(progress);
+                }
+            );
 
             await this.asr.start();
 
@@ -769,6 +826,9 @@ class ASRDemoUI {
 
     showLoading(text) {
         this.loadingText.textContent = text;
+        this.loadingProgressFill.style.width = '0%';
+        this.loadingProgressText.textContent = '0 / 7 models';
+        this.loadingDetails.textContent = '';
         this.loadingOverlay.classList.remove('hidden');
     }
 
@@ -776,6 +836,25 @@ class ASRDemoUI {
         this.loadingOverlay.classList.add('hidden');
     }
 
+    updateLoadingProgress(progress) {
+        const { completedModels, totalModels, currentModel, currentProgress } = progress;
+
+        // Update overall progress bar
+        const overallPercent = (completedModels / totalModels) * 100;
+        this.loadingProgressFill.style.width = `${overallPercent}%`;
+        this.loadingProgressText.textContent = `${completedModels} / ${totalModels} models`;
+
+        // Update details with current model and byte progress
+        if (currentModel && currentProgress.total > 0) {
+            const loadedMB = (currentProgress.loaded / (1024 * 1024)).toFixed(1);
+            const totalMB = (currentProgress.total / (1024 * 1024)).toFixed(1);
+            const percent = Math.round((currentProgress.loaded / currentProgress.total) * 100);
+            this.loadingDetails.textContent = `${currentModel}: ${loadedMB} / ${totalMB} MB (${percent}%)`;
+        } else if (currentModel) {
+            this.loadingDetails.textContent = `Loading ${currentModel}...`;
+        }
+    }
+
     showError(message) {
         this.errorMessage.textContent = message;
         this.errorMessage.classList.add('visible');