Handle the separate codebook step

audio-model.js

@@ -397,7 +397,7 @@ export class AudioModel {
       };
 
       // Helper to load ONNX model with external data
-      const loadOnnxWithExternalData = async (name, progress, quantSuffix = null) => {
+      const loadOnnxWithExternalData = async (name, progress, quantSuffix = null, epOverride = null, extraOptions = {}) => {
         const suffix = quantSuffix ? `_${quantSuffix}` : '';
         const fileName = `${name}${suffix}`;
         report('loading', progress, `${fileName}.onnx`);
@@ -405,9 +405,10 @@ export class AudioModel {
         const onnxPath = `${modelPath}/onnx/${fileName}.onnx`;
         const fetchOptions = { mode: 'cors', credentials: 'omit' };
 
-        console.log(`Loading ${fileName}...`);
+        const ep = epOverride || executionProviders;
+        console.log(`Loading ${fileName} (EP: ${JSON.stringify(ep)})...`);
 
-        const sessionOptions = { executionProviders };
+        const sessionOptions = { executionProviders: ep, ...extraOptions };
 
         const onnxResponse = await fetchWithCache(onnxPath, fetchOptions);
         if (!onnxResponse.ok) {
@@ -464,7 +465,20 @@ export class AudioModel {
       };
 
       // Load decoder
-      this.decoderSession = await loadOnnxWithExternalData('decoder', 10, quantConfig.decoder);
+      // On WebGPU: keep KV cache outputs on GPU to avoid GPU→CPU→GPU roundtrips between steps
+      const decoderOpts = device === 'webgpu' ? (() => {
+        const loc = {};
+        for (let i = 0; i < this.layerTypes.length; i++) {
+          if (this.layerTypes[i] === 'conv') {
+            loc[`present_conv.${i}`] = 'gpu-buffer';
+          } else {
+            loc[`present.${i}.key`] = 'gpu-buffer';
+            loc[`present.${i}.value`] = 'gpu-buffer';
+          }
+        }
+        return { preferredOutputLocation: loc };
+      })() : {};
+      this.decoderSession = await loadOnnxWithExternalData('decoder', 10, quantConfig.decoder, null, decoderOpts);
 
       // Load embed_tokens weight for text embedding lookup
       report('loading', 30, 'embed_tokens');
@@ -490,9 +504,15 @@ export class AudioModel {
         console.warn('Audio detokenizer not available:', e);
       }
 
-      // Load vocoder (for TTS)
+      // Load vocoder/depthformer (for TTS) — per-step model (8 calls per frame)
+      // On WebGPU: keep KV cache on GPU to avoid GPU→CPU→GPU roundtrips between steps
       try {
-        this.vocoderSession = await loadOnnxWithExternalData('vocoder_depthformer', 95, quantConfig.vocoder);
+        const vocoderOpts = device === 'webgpu'
+          ? { preferredOutputLocation: { new_keys: 'gpu-buffer', new_values: 'gpu-buffer', depth_slices: 'gpu-buffer' } }
+          : {};
+        this.vocoderSession = await loadOnnxWithExternalData(
+          'vocoder_depthformer', 95, quantConfig.vocoder, null, vocoderOpts,
+        );
       } catch (e) {
         console.warn('Vocoder not available:', e);
       }
@@ -934,129 +954,103 @@ export class AudioModel {
     return '[Text generation requires full embedding support - model loaded successfully]';
   }
 
-  /**
-   * Initialize reusable vocoder tensors to reduce allocation overhead
-   */
   _initVocoderCache() {
     if (this._vocoderCache) return;
 
-    const numLayers = 6;
-    const numKvHeads = 8;
-    const headDim = 32;
-
-    // Pre-allocate data arrays
+    const vocabSize = 2049;
     const stepIdxData = new BigInt64Array(1);
     const prevTokenData = new BigInt64Array(1);
+    const seqlensKData = new Int32Array(1);
+    const totalSeqLenData = new Int32Array(1);
 
-    // Pre-allocate tensors that can be reused
     this._vocoderCache = {
-      hiddenTensor: null,  // Created per-call since hiddenState changes
       stepIdxData,
       prevTokenData,
-      // Pre-create reusable tensors (ONNX Runtime reads from the data array)
+      seqlensKData,
+      totalSeqLenData,
       stepIdxTensor: new ort.Tensor('int64', stepIdxData, []),
       prevTokenTensor: new ort.Tensor('int64', prevTokenData, [1]),
-      emptyKeysData: new Float32Array(0),
-      emptyValuesData: new Float32Array(0),
-      // Reusable sampling arrays
-      scaledLogits: new Float32Array(2049),  // codebook vocab size
-      indices: new Uint16Array(2049),  // Use typed array for faster reset
-      probs: new Float32Array(64),  // top-k size
+      seqlensKTensor: new ort.Tensor('int32', seqlensKData, [1]),
+      totalSeqLenTensor: new ort.Tensor('int32', totalSeqLenData, []),
+      emptyData: new Float32Array(0),
+      // Pre-allocated sampling arrays
+      scaledLogits: new Float32Array(vocabSize),
+      indices: new Uint16Array(vocabSize),
+      probs: new Float32Array(64),
     };
-
-    // Initialize indices
-    for (let i = 0; i < 2049; i++) {
-      this._vocoderCache.indices[i] = i;
-    }
   }
 
   /**
-   * Sample audio codes using vocoder depthformer
-   * Optimized to reduce tensor creation overhead
+   * Sample audio codes using per-step depthformer (8 session.run calls).
+   * Uses GroupQueryAttention with BNSH KV cache format.
    * @param {Float32Array} hiddenState - [hidden_size] hidden state
    * @param {number} temperature - Sampling temperature
-   * @param {number} topK - Top-k sampling
+   * @param {number} topK - Top-k sampling (0 = greedy)
    * @returns {number[]} - 8 codebook values
    */
   async sampleAudioCodes(hiddenState, temperature = 0.8, topK = 64) {
-    if (!this.vocoderSession) {
-      throw new Error('Vocoder not loaded');
-    }
-
-    // Initialize cache on first call
     this._initVocoderCache();
     const cache = this._vocoderCache;
 
     const numCodebooks = 8;
     const numLayers = 6;
-    const numKvHeads = 8;
+    const numKVHeads = 8;
     const headDim = 32;
+    const vocabSize = 2049;
 
-    const codes = [];
-    let prevToken = 0;
-
-    // Create hidden state tensor (must be new since data changes)
     const hiddenTensor = new ort.Tensor('float32', hiddenState, [1, this.hiddenSize]);
+    // BNSH format: [layers, batch, heads, seq_len, head_dim]
+    let pastKeys = new ort.Tensor('float32', cache.emptyData, [numLayers, 1, numKVHeads, 0, headDim]);
+    let pastValues = new ort.Tensor('float32', cache.emptyData, [numLayers, 1, numKVHeads, 0, headDim]);
+    let depthSlices = new ort.Tensor('float32', new Float32Array(numCodebooks * 1024), [1, numCodebooks, 1024]);
 
-    // Initialize empty KV cache
-    let pastKeys = new ort.Tensor(
-      'float32',
-      cache.emptyKeysData,
-      [numLayers, 1, 0, numKvHeads, headDim]
-    );
-    let pastValues = new ort.Tensor(
-      'float32',
-      cache.emptyValuesData,
-      [numLayers, 1, 0, numKvHeads, headDim]
-    );
-
-    // Reuse step_idx and prev_token tensors by updating their data
-    cache.stepIdxData[0] = 0n;
-    cache.prevTokenData[0] = 0n;
+    const codes = [];
+    let prevToken = 0;
 
-    for (let i = 0; i < numCodebooks; i++) {
-      // Update mutable tensor data (tensor objects reuse the underlying data arrays)
-      cache.stepIdxData[0] = BigInt(i);
+    for (let step = 0; step < numCodebooks; step++) {
+      cache.stepIdxData[0] = BigInt(step);
       cache.prevTokenData[0] = BigInt(prevToken);
+      cache.seqlensKData[0] = step;
+      cache.totalSeqLenData[0] = step + 1;
 
-      const feeds = {
+      const outputs = await this.vocoderSession.run({
         hidden_states: hiddenTensor,
+        depth_slices_in: depthSlices,
         step_idx: cache.stepIdxTensor,
         prev_token: cache.prevTokenTensor,
         past_keys: pastKeys,
         past_values: pastValues,
-      };
+        seqlens_k: cache.seqlensKTensor,
+        total_seq_len: cache.totalSeqLenTensor,
+      });
+
+      if (step === 0) {
+        depthSlices = outputs.depth_slices;
+      }
+
+      pastKeys = outputs.new_keys;
+      pastValues = outputs.new_values;
 
-      const outputs = await this.vocoderSession.run(feeds);
       const logits = outputs.logits.data;
-      const vocabSize = logits.length;
 
-      // Sample with temperature and top-k (reusing cached arrays)
       let token;
-      if (temperature <= 0) {
-        // Greedy
+      if (temperature <= 0 || topK <= 1) {
         token = 0;
         let maxVal = logits[0];
         for (let j = 1; j < vocabSize; j++) {
-          if (logits[j] > maxVal) {
-            maxVal = logits[j];
-            token = j;
-          }
+          if (logits[j] > maxVal) { maxVal = logits[j]; token = j; }
         }
       } else {
-        // Top-k sampling with reused arrays
         const scaledLogits = cache.scaledLogits;
         const indices = cache.indices;
         const probs = cache.probs;
 
-        // Scale logits by temperature and find top-k in single pass
-        // Use partial selection sort (O(k*n) which is fast for small k)
         for (let j = 0; j < vocabSize; j++) {
           scaledLogits[j] = logits[j] / temperature;
           indices[j] = j;
         }
 
-        // Partial sort to get top-k
+        // Partial selection sort for top-k: O(k*n) vs O(n log n) full sort
         for (let j = 0; j < topK; j++) {
           let maxIdx = j;
           for (let k = j + 1; k < vocabSize; k++) {
@@ -1064,7 +1058,6 @@ export class AudioModel {
               maxIdx = k;
             }
           }
-          // Swap
           const tmp = indices[j];
           indices[j] = indices[maxIdx];
           indices[maxIdx] = tmp;
@@ -1081,25 +1074,18 @@ export class AudioModel {
           probs[j] /= sumExp;
         }
 
-        // Sample
+        // Sample from cumulative distribution
         const r = Math.random();
         let cumsum = 0;
-        token = indices[topK - 1];  // Default to last
+        token = indices[topK - 1];
         for (let j = 0; j < topK; j++) {
           cumsum += probs[j];
-          if (r < cumsum) {
-            token = indices[j];
-            break;
-          }
+          if (r < cumsum) { token = indices[j]; break; }
         }
       }
 
       codes.push(token);
       prevToken = token;
-
-      // Update KV cache
-      pastKeys = outputs.new_keys;
-      pastValues = outputs.new_values;
     }
 
     return codes;
@@ -1304,12 +1290,15 @@ export class AudioModel {
       throw new Error('Vocoder not loaded - required for interleaved mode');
     }
 
-    // Timing accumulators
+    // Timing accumulators (names match liquid-audio architecture):
+    // - lfmText/lfmAudio: self.lfm() decoder calls (text vs audio steps)
+    // - depthformer: self._sample_audio_frame() — depth_linear + 8× depthformer
+    // - audioEmbed: self.audio_embedding(...).sum() — feedback embedding
     let timeAudioEncode = 0;
     let timePrefill = 0;
-    let timeTextDecode = 0;
-    let timeAudioDecode = 0;
-    let timeVocoder = 0;
+    let timeLfmText = 0;
+    let timeLfmAudio = 0;
+    let timeDepthformer = 0;
     let timeAudioEmbed = 0;
 
     // 1. Compute mel spectrogram and encode audio
@@ -1415,36 +1404,37 @@ export class AudioModel {
 
     const startTime = performance.now();
 
-    for (let step = 0; step < maxNewTokens; step++) {
+    log(`Generation loop: max ${maxNewTokens} steps, starting in TEXT mode`);
+
+    let step = 0;
+    for (; step < maxNewTokens; step++) {
       modalityLeft--;
 
       if (inAudioMode) {
-        // Generate audio frame using depthformer
+        // === AUDIO STEP: extract hidden_states → depthformer → 8 codebook tokens ===
         const hiddenData = hiddenStates.data;
         const seqLen = hiddenStates.dims[1];
         const lastHidden = hiddenData.slice((seqLen - 1) * hiddenSize, seqLen * hiddenSize);
 
         tStep = performance.now();
         const frameCodes = await this.sampleAudioCodes(lastHidden, audioTemperature, audioTopK);
-        timeVocoder += performance.now() - tStep;
+        timeDepthformer += performance.now() - tStep;
 
         // Switch back to text after N audio frames (if text not done)
         if (modalityLeft <= 0 && !textDone) {
+          log(`→ AUDIO→TEXT (after ${INTERLEAVED_N_AUDIO} audio frames, ${audioCodes.length} total)`);
           inAudioMode = false;
           modalityLeft = INTERLEAVED_N_TEXT;
         }
 
         // Check for end of audio - first codebook == 2048 (matching liquid-audio)
         if (frameCodes[0] === END_OF_AUDIO_TOKEN) {
-          log(`End of audio at step ${step}`);
-          // Set all codes to 2048 (matching liquid-audio)
+          log(`→ END_OF_AUDIO at step ${step} (${audioCodes.length} frames collected)`);
           for (let i = 0; i < NUM_CODEBOOKS; i++) {
             frameCodes[i] = END_OF_AUDIO_TOKEN;
           }
           inAudioMode = false;
-          // Don't save this frame, but still feed it back
         } else {
-          // Save valid frame (clamped to 0-2047)
           const clampedFrame = frameCodes.map(c => Math.min(c, 2047));
           audioCodes.push(clampedFrame);
 
@@ -1453,16 +1443,15 @@ export class AudioModel {
           }
 
           if (audioCodes.length % 50 === 0) {
-            log(`Generated ${audioCodes.length} audio frames`);
+            log(`  Audio frames: ${audioCodes.length}`);
           }
         }
 
-        // Get embeddings for next step (always feed back, even for 2048 frames)
+        // === FEEDBACK: embed 8 codes (summed) → feed back to LFM decoder ===
         tStep = performance.now();
         const feedCodes = frameCodes.map(c => c === END_OF_AUDIO_TOKEN ? END_OF_AUDIO_TOKEN : Math.min(c, 2047));
         const audioTokens = feedCodes.map((code, idx) => idx * CODEBOOK_VOCAB + code);
 
-        // Get summed embeddings for all 8 codebooks
         const summedEmbeds = await this.getAudioEmbedding(audioTokens);
         timeAudioEmbed += performance.now() - tStep;
 
@@ -1471,14 +1460,13 @@ export class AudioModel {
         const nextMask = new ort.Tensor('int64', new BigInt64Array(currentLen).fill(1n), [1, currentLen]);
         tStep = performance.now();
         ({ logits, hiddenStates, outputs } = await this.runDecoder(nextEmbeds, nextMask, this.cache));
-        timeAudioDecode += performance.now() - tStep;
+        timeLfmAudio += performance.now() - tStep;
         this.updateCache(this.cache, outputs);
 
       } else {
-        // Generate text token
+        // === TEXT STEP: logits → sample text token ===
         const logitsData = logits.data;
         const seqLen = logits.dims[1];
-        // Get logits for last position - shape is [1, seq_len, vocab_size]
         const lastLogits = new Float32Array(this.vocabSize);
         const offset = (seqLen - 1) * this.vocabSize;
         for (let i = 0; i < this.vocabSize; i++) {
@@ -1488,18 +1476,19 @@ export class AudioModel {
 
         // Check for end of turn
         if (token === this.tokenizer.eos_token_id || token === SPECIAL_TOKENS.IM_END) {
-          log(`End of turn at step ${step}`);
+          log(`→ END_OF_TURN at step ${step} (${textTokens.length} text tokens, ${audioCodes.length} audio frames)`);
           break;
         }
 
         // Check for <|text_end|> token (130)
         if (token === SPECIAL_TOKENS.TEXT_END) {
-          log(`Text end at step ${step}`);
+          log(`→ TEXT_END at step ${step}: audio-only phase begins`);
           textDone = true;
         }
 
         // Switch to audio after N text tokens OR text_end
         if (modalityLeft <= 0 || textDone) {
+          log(`→ TEXT→AUDIO${textDone ? ' (text_done)' : ''} at step ${step}`);
           inAudioMode = true;
           modalityLeft = INTERLEAVED_N_AUDIO;
         }
@@ -1511,18 +1500,19 @@ export class AudioModel {
           onToken(decodedText, token);
         }
 
-        // Get embedding for next step
+        // === FEEDBACK: embed text token → feed back to LFM decoder ===
         const nextEmbeds = this.getTextEmbeddings([token]);
         currentLen++;
         const nextMask = new ort.Tensor('int64', new BigInt64Array(currentLen).fill(1n), [1, currentLen]);
         tStep = performance.now();
         ({ logits, hiddenStates, outputs } = await this.runDecoder(nextEmbeds, nextMask, this.cache));
-        timeTextDecode += performance.now() - tStep;
+        timeLfmText += performance.now() - tStep;
         this.updateCache(this.cache, outputs);
       }
     }
 
     // 5. Feed <|im_end|> token to close assistant turn in cache
+
     const imEndEmbeds = this.getTextEmbeddings([SPECIAL_TOKENS.IM_END]);
     currentLen++;
     const finalMask = new ort.Tensor('int64', new BigInt64Array(currentLen).fill(1n), [1, currentLen]);
@@ -1533,14 +1523,11 @@ export class AudioModel {
     // Decode with skip_special_tokens to clean up special tokens like <|text_end|>
     const text = this.tokenizer.decode(textTokens, { skip_special_tokens: true });
 
-    // Print timing summary
-    log(`=== Summary ===`);
-    log(`  Mel: ${timeMel.toFixed(0)}ms, AudioEnc: ${timeAudioEncode.toFixed(0)}ms, Prefill: ${timePrefill.toFixed(0)}ms`);
-    log(`  TextDec: ${timeTextDecode.toFixed(0)}ms (${textTokens.length} tok), AudioDec: ${timeAudioDecode.toFixed(0)}ms`);
-    log(`  Vocoder: ${timeVocoder.toFixed(0)}ms, AudioEmbed: ${timeAudioEmbed.toFixed(0)}ms`);
-    log(`Output: ${textTokens.length} text tokens, ${audioCodes.length} audio frames`);
-    log(`Text: "${text}"`);
-    log(`Cache seq_len: ${this.cacheSeqLen}`);
+    const totalGenTime = performance.now() - startTime;
+
+    log(`Done: ${step} steps in ${totalGenTime.toFixed(0)}ms | ${textTokens.length} text tokens, ${audioCodes.length} audio frames (~${(audioCodes.length / 75).toFixed(1)}s audio)`);
+    log(`Timing: mel=${timeMel.toFixed(0)}ms, audioEnc=${timeAudioEncode.toFixed(0)}ms, prefill=${timePrefill.toFixed(0)}ms, lfmText=${timeLfmText.toFixed(0)}ms, lfmAudio=${timeLfmAudio.toFixed(0)}ms, depthformer=${timeDepthformer.toFixed(0)}ms, audioEmbed=${timeAudioEmbed.toFixed(0)}ms`);
+    log(`Text: "${text}" | cache_seq_len=${this.cacheSeqLen}`);
 
     return { text, audioCodes };
   }
@@ -1581,11 +1568,14 @@ export class AudioModel {
       throw new Error('Vocoder not loaded - required for interleaved mode');
     }
 
-    // Timing accumulators
+    // Timing accumulators (names match liquid-audio architecture):
+    // - lfmText/lfmAudio: self.lfm() decoder calls (text vs audio steps)
+    // - depthformer: self._sample_audio_frame() — depth_linear + 8× depthformer
+    // - audioEmbed: self.audio_embedding(...).sum() — feedback embedding
     let timePrefill = 0;
-    let timeTextDecode = 0;
-    let timeAudioDecode = 0;
-    let timeVocoder = 0;
+    let timeLfmText = 0;
+    let timeLfmAudio = 0;
+    let timeDepthformer = 0;
     let timeAudioEmbed = 0;
     let tStep;
 
@@ -1650,17 +1640,18 @@ export class AudioModel {
 
         tStep = performance.now();
         const frameCodes = await this.sampleAudioCodes(lastHidden, audioTemperature, audioTopK);
-        timeVocoder += performance.now() - tStep;
+        timeDepthformer += performance.now() - tStep;
 
         // Switch back to text after N audio frames (if text not done)
         if (modalityLeft <= 0 && !textDone) {
+          log(`→ AUDIO→TEXT (after ${INTERLEAVED_N_AUDIO} audio frames, ${audioCodes.length} total)`);
           inAudioMode = false;
           modalityLeft = INTERLEAVED_N_TEXT;
         }
 
         // Check for end of audio
         if (frameCodes[0] === END_OF_AUDIO_TOKEN) {
-          log(`End of audio at step ${step}`);
+          log(`→ END_OF_AUDIO at step ${step} (${audioCodes.length} frames collected)`);
           for (let i = 0; i < NUM_CODEBOOKS; i++) {
             frameCodes[i] = END_OF_AUDIO_TOKEN;
           }
@@ -1674,7 +1665,7 @@ export class AudioModel {
           }
 
           if (audioCodes.length % 50 === 0) {
-            log(`Generated ${audioCodes.length} audio frames`);
+            log(`  Audio frames: ${audioCodes.length}`);
           }
         }
 
@@ -1690,7 +1681,7 @@ export class AudioModel {
         const nextMask = new ort.Tensor('int64', new BigInt64Array(currentLen).fill(1n), [1, currentLen]);
         tStep = performance.now();
         ({ logits, hiddenStates, outputs } = await this.runDecoder(nextEmbeds, nextMask, this.cache));
-        timeAudioDecode += performance.now() - tStep;
+        timeLfmAudio += performance.now() - tStep;
         this.updateCache(this.cache, outputs);
 
       } else {
@@ -1706,18 +1697,19 @@ export class AudioModel {
 
         // Check for end of turn
         if (token === this.tokenizer.eos_token_id || token === SPECIAL_TOKENS.IM_END) {
-          log(`End of turn at step ${step}`);
+          log(`→ END_OF_TURN at step ${step} (${textTokens.length} text tokens, ${audioCodes.length} audio frames)`);
           break;
         }
 
         // Check for <|text_end|> token
         if (token === SPECIAL_TOKENS.TEXT_END) {
-          log(`Text end at step ${step}`);
+          log(`→ TEXT_END at step ${step}: audio-only phase begins`);
           textDone = true;
         }
 
         // Switch to audio after N text tokens OR text_end
         if (modalityLeft <= 0 || textDone) {
+          log(`→ TEXT→AUDIO${textDone ? ' (text_done)' : ''} at step ${step}`);
           inAudioMode = true;
           modalityLeft = INTERLEAVED_N_AUDIO;
         }
@@ -1735,7 +1727,7 @@ export class AudioModel {
         const nextMask = new ort.Tensor('int64', new BigInt64Array(currentLen).fill(1n), [1, currentLen]);
         tStep = performance.now();
         ({ logits, hiddenStates, outputs } = await this.runDecoder(nextEmbeds, nextMask, this.cache));
-        timeTextDecode += performance.now() - tStep;
+        timeLfmText += performance.now() - tStep;
         this.updateCache(this.cache, outputs);
       }
     }
@@ -1750,13 +1742,9 @@ export class AudioModel {
 
     const text = this.tokenizer.decode(textTokens, { skip_special_tokens: true });
 
-    log(`=== Summary ===`);
-    log(`  Prefill: ${timePrefill.toFixed(0)}ms`);
-    log(`  TextDec: ${timeTextDecode.toFixed(0)}ms (${textTokens.length} tok), AudioDec: ${timeAudioDecode.toFixed(0)}ms`);
-    log(`  Vocoder: ${timeVocoder.toFixed(0)}ms, AudioEmbed: ${timeAudioEmbed.toFixed(0)}ms`);
-    log(`Output: ${textTokens.length} text tokens, ${audioCodes.length} audio frames`);
-    log(`Text: "${text}"`);
-    log(`Cache seq_len: ${this.cacheSeqLen}`);
+    log(`Done: ${textTokens.length} text tokens, ${audioCodes.length} audio frames (~${(audioCodes.length / 75).toFixed(1)}s audio)`);
+    log(`Timing: prefill=${timePrefill.toFixed(0)}ms, lfmText=${timeLfmText.toFixed(0)}ms, lfmAudio=${timeLfmAudio.toFixed(0)}ms, depthformer=${timeDepthformer.toFixed(0)}ms, audioEmbed=${timeAudioEmbed.toFixed(0)}ms`);
+    log(`Text: "${text}" | cache_seq_len=${this.cacheSeqLen}`);
 
     return { text, audioCodes };
   }
@@ -1886,7 +1874,8 @@ export class AudioModel {
     }
 
     const decodeStart = performance.now();
-    log(`Decoding ${audioCodes.length} audio frames...`);
+
+    log(`Audio decode: ${audioCodes.length} frames → waveform`);
 
     // ISTFT parameters (fixed for this model)
     const nFft = 1280;
@@ -1894,7 +1883,7 @@ export class AudioModel {
     const winLength = 1280;
     const nFftBins = nFft / 2 + 1;
 
-    // Stack codes: [T, 8] -> [8, T] and add batch -> [1, 8, T]
+    // Transpose codes [T, 8] → [1, 8, T] for ONNX input
     const T = audioCodes.length;
     const codesTransposed = new BigInt64Array(8 * T);
     for (let t = 0; t < T; t++) {
@@ -1903,18 +1892,18 @@ export class AudioModel {
       }
     }
 
-    // Run detokenizer: [1, 8, T] -> [1, T, 1282]
+    // Run detokenizer ONNX: [1, 8, T] → [1, 6T, 1282]
     const codesTensor = new ort.Tensor('int64', codesTransposed, [1, 8, T]);
     const detokStart = performance.now();
     const detokOutputs = await this.audioDetokenizerSession.run({ audio_codes: codesTensor });
     const stftFeatures = detokOutputs.stft_features;
-    log(`Detokenizer: ${(performance.now() - detokStart).toFixed(0)}ms, STFT frames: ${stftFeatures.dims[1]}`);
+    const actualT = stftFeatures.dims[1];
+    const detokEnd = performance.now();
+    log(`  Detokenizer: [1,8,${T}] → [1,${actualT},1282] in ${(detokEnd - detokStart).toFixed(0)}ms`);
 
-    // Get raw data - shape is [1, T, 1282], we need to skip batch dimension
+    // Split into magnitude + angle → complex spectrogram
     const stftData = stftFeatures.data;
-    const actualT = stftFeatures.dims[1];
 
-    // Convert to complex STFT: [log_magnitude | angle] -> complex
     const complexStft = new Array(nFftBins);
     for (let f = 0; f < nFftBins; f++) {
       complexStft[f] = new Array(actualT);
@@ -1922,23 +1911,23 @@ export class AudioModel {
         const logMag = stftData[t * 1282 + f];
         const angle = stftData[t * 1282 + nFftBins + f];
         const mag = Math.exp(logMag);
-        // Store as [real, imag]
         complexStft[f][t] = [mag * Math.cos(angle), mag * Math.sin(angle)];
       }
     }
 
-    // ISTFT with 'same' padding
+    // ISTFT (inverse Short-Time Fourier Transform) → waveform
+    const pad = (winLength - hopLength) / 2;
     const istftStart = performance.now();
     const waveform = this.istftSamePadding(complexStft, nFft, hopLength, winLength, actualT);
-    log(`ISTFT: ${(performance.now() - istftStart).toFixed(0)}ms`);
+    const istftEnd = performance.now();
+    log(`  ISTFT: ${actualT} frames → ${waveform.length} samples in ${(istftEnd - istftStart).toFixed(0)}ms`);
 
-    // Find max/min without spread operator (avoid stack overflow on large arrays)
+    // Find max/min
     let waveMax = -Infinity, waveMin = Infinity;
     for (let i = 0; i < waveform.length; i++) {
       if (waveform[i] > waveMax) waveMax = waveform[i];
       if (waveform[i] < waveMin) waveMin = waveform[i];
     }
-    log('ISTFT output - length:', waveform.length, 'max:', waveMax.toFixed(4), 'min:', waveMin.toFixed(4));
 
     // Check for invalid values
     if (isNaN(waveMax) || isNaN(waveMin) || !isFinite(waveMax) || !isFinite(waveMin)) {
@@ -1946,7 +1935,7 @@ export class AudioModel {
       return new Float32Array(0);
     }
 
-    // Normalize to [-1, 1]
+    // Normalize waveform to [-0.9, 0.9]
     let maxVal = Math.max(Math.abs(waveMax), Math.abs(waveMin));
     if (maxVal > 0) {
       for (let i = 0; i < waveform.length; i++) {
@@ -1956,7 +1945,8 @@ export class AudioModel {
       console.warn('ISTFT produced all-zero waveform');
     }
 
-    log(`Decoded audio: ${waveform.length} samples (${(waveform.length / 24000).toFixed(2)}s)`);
+    const totalDecodeTime = performance.now() - decodeStart;
+    log(`Audio decode complete: ${totalDecodeTime.toFixed(0)}ms total (detok=${(detokEnd - detokStart).toFixed(0)}ms, istft=${(istftEnd - istftStart).toFixed(0)}ms) → ${waveform.length} samples, ${(waveform.length / 24000).toFixed(2)}s @ 24kHz`);
     return waveform;
   }
 

main.js

@@ -6,13 +6,13 @@
 
 import { AudioModel, loadAudioFile, clearModelCache, getCacheInfo } from './audio-model.js';
 
-// HuggingFace model URL
-const MODEL_URL = 'https://huggingface.co/LiquidAI/LFM2.5-Audio-1.5B-ONNX/resolve/main';
+// Model path - local directory
+const MODEL_PATH = './LFM2.5-Audio-1.5B-ONNX';
 
 // Model configurations
 const MODELS = {
   'LFM2.5-Audio-1.5B-Q4': {
-    path: MODEL_URL,
+    path: MODEL_PATH,
     label: 'LFM2.5-Audio-1.5B Q4 (~1.6 GB)',
     quantization: {
       decoder: 'q4',