correct the input

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, epOverride = null, extraOptions = {}) => {
+      const loadOnnxWithExternalData = async (name, progress, quantSuffix = null, extraOptions = {}) => {
         const suffix = quantSuffix ? `_${quantSuffix}` : '';
         const fileName = `${name}${suffix}`;
         report('loading', progress, `${fileName}.onnx`);
@@ -405,10 +405,9 @@ export class AudioModel {
         const onnxPath = `${modelPath}/onnx/${fileName}.onnx`;
         const fetchOptions = { mode: 'cors', credentials: 'omit' };
 
-        const ep = epOverride || executionProviders;
-        console.log(`Loading ${fileName} (EP: ${JSON.stringify(ep)})...`);
+        console.log(`Loading ${fileName}...`);
 
-        const sessionOptions = { executionProviders: ep, ...extraOptions };
+        const sessionOptions = { executionProviders, ...extraOptions };
 
         const onnxResponse = await fetchWithCache(onnxPath, fetchOptions);
         if (!onnxResponse.ok) {
@@ -478,7 +477,7 @@ export class AudioModel {
         }
         return { preferredOutputLocation: loc };
       })() : {};
-      this.decoderSession = await loadOnnxWithExternalData('decoder', 10, quantConfig.decoder, null, decoderOpts);
+      this.decoderSession = await loadOnnxWithExternalData('decoder', 10, quantConfig.decoder, decoderOpts);
 
       // Load embed_tokens weight for text embedding lookup
       report('loading', 30, 'embed_tokens');
@@ -504,15 +503,13 @@ export class AudioModel {
         console.warn('Audio detokenizer not available:', e);
       }
 
-      // Load vocoder/depthformer (for TTS) — per-step model (8 calls per frame)
+      // Load vocoder (for TTS)
       // On WebGPU: keep KV cache on GPU to avoid GPU→CPU→GPU roundtrips between steps
       try {
         const vocoderOpts = device === 'webgpu'
           ? { preferredOutputLocation: { new_keys: 'gpu-buffer', new_values: 'gpu-buffer' } }
           : {};
-        this.vocoderSession = await loadOnnxWithExternalData(
-          'vocoder_depthformer', 95, quantConfig.vocoder, null, vocoderOpts,
-        );
+        this.vocoderSession = await loadOnnxWithExternalData('vocoder_depthformer', 95, quantConfig.vocoder, vocoderOpts);
       } catch (e) {
         console.warn('Vocoder not available:', e);
       }
@@ -954,103 +951,129 @@ 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 vocabSize = 2049;
+    const numLayers = 6;
+    const numKvHeads = 8;
+    const headDim = 32;
+
+    // Pre-allocate data arrays
     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,
-      seqlensKData,
-      totalSeqLenData,
+      // Pre-create reusable tensors (ONNX Runtime reads from the data array)
       stepIdxTensor: new ort.Tensor('int64', stepIdxData, []),
       prevTokenTensor: new ort.Tensor('int64', prevTokenData, [1]),
-      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),
+      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
     };
+
+    // Initialize indices
+    for (let i = 0; i < 2049; i++) {
+      this._vocoderCache.indices[i] = i;
+    }
   }
 
   /**
-   * Sample audio codes using per-step depthformer (8 session.run calls).
-   * Uses GroupQueryAttention with BNSH KV cache format.
+   * Sample audio codes using vocoder depthformer
+   * Optimized to reduce tensor creation overhead
    * @param {Float32Array} hiddenState - [hidden_size] hidden state
    * @param {number} temperature - Sampling temperature
-   * @param {number} topK - Top-k sampling (0 = greedy)
+   * @param {number} topK - Top-k sampling
    * @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 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]);
 
     const codes = [];
     let prevToken = 0;
 
-    for (let step = 0; step < numCodebooks; step++) {
-      cache.stepIdxData[0] = BigInt(step);
+    // Create hidden state tensor (must be new since data changes)
+    const hiddenTensor = new ort.Tensor('float32', hiddenState, [1, this.hiddenSize]);
+
+    // 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;
+
+    for (let i = 0; i < numCodebooks; i++) {
+      // Update mutable tensor data (tensor objects reuse the underlying data arrays)
+      cache.stepIdxData[0] = BigInt(i);
       cache.prevTokenData[0] = BigInt(prevToken);
-      cache.seqlensKData[0] = step;
-      cache.totalSeqLenData[0] = step + 1;
 
-      const outputs = await this.vocoderSession.run({
+      const feeds = {
         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 || topK <= 1) {
+      if (temperature <= 0) {
+        // Greedy
         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 selection sort for top-k: O(k*n) vs O(n log n) full sort
+        // Partial sort to get top-k
         for (let j = 0; j < topK; j++) {
           let maxIdx = j;
           for (let k = j + 1; k < vocabSize; k++) {
@@ -1058,6 +1081,7 @@ export class AudioModel {
               maxIdx = k;
             }
           }
+          // Swap
           const tmp = indices[j];
           indices[j] = indices[maxIdx];
           indices[maxIdx] = tmp;
@@ -1074,18 +1098,25 @@ export class AudioModel {
           probs[j] /= sumExp;
         }
 
-        // Sample from cumulative distribution
+        // Sample
         const r = Math.random();
         let cumsum = 0;
-        token = indices[topK - 1];
+        token = indices[topK - 1];  // Default to last
         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;
@@ -1290,15 +1321,12 @@ export class AudioModel {
       throw new Error('Vocoder not loaded - required for interleaved mode');
     }
 
-    // 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
+    // Timing accumulators
     let timeAudioEncode = 0;
     let timePrefill = 0;
-    let timeLfmText = 0;
-    let timeLfmAudio = 0;
-    let timeDepthformer = 0;
+    let timeTextDecode = 0;
+    let timeAudioDecode = 0;
+    let timeVocoder = 0;
     let timeAudioEmbed = 0;
 
     // 1. Compute mel spectrogram and encode audio
@@ -1404,37 +1432,36 @@ export class AudioModel {
 
     const startTime = performance.now();
 
-    log(`Generation loop: max ${maxNewTokens} steps, starting in TEXT mode`);
-
-    let step = 0;
-    for (; step < maxNewTokens; step++) {
+    for (let step = 0; step < maxNewTokens; step++) {
       modalityLeft--;
 
       if (inAudioMode) {
-        // === AUDIO STEP: extract hidden_states → depthformer → 8 codebook tokens ===
+        // Generate audio frame using depthformer
         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);
-        timeDepthformer += performance.now() - tStep;
+        timeVocoder += 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} (${audioCodes.length} frames collected)`);
+          log(`End of audio at step ${step}`);
+          // Set all codes to 2048 (matching liquid-audio)
           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);
 
@@ -1443,15 +1470,16 @@ export class AudioModel {
           }
 
           if (audioCodes.length % 50 === 0) {
-            log(`  Audio frames: ${audioCodes.length}`);
+            log(`Generated ${audioCodes.length} audio frames`);
           }
         }
 
-        // === FEEDBACK: embed 8 codes (summed) → feed back to LFM decoder ===
+        // Get embeddings for next step (always feed back, even for 2048 frames)
         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;
 
@@ -1460,13 +1488,14 @@ 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));
-        timeLfmAudio += performance.now() - tStep;
+        timeAudioDecode += performance.now() - tStep;
         this.updateCache(this.cache, outputs);
 
       } else {
-        // === TEXT STEP: logits → sample text token ===
+        // Generate 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++) {
@@ -1476,19 +1505,18 @@ 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} (${textTokens.length} text tokens, ${audioCodes.length} audio frames)`);
+          log(`End of turn at step ${step}`);
           break;
         }
 
         // Check for <|text_end|> token (130)
         if (token === SPECIAL_TOKENS.TEXT_END) {
-          log(`→ TEXT_END at step ${step}: audio-only phase begins`);
+          log(`Text end at step ${step}`);
           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;
         }
@@ -1500,19 +1528,18 @@ export class AudioModel {
           onToken(decodedText, token);
         }
 
-        // === FEEDBACK: embed text token → feed back to LFM decoder ===
+        // Get embedding for next step
         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));
-        timeLfmText += performance.now() - tStep;
+        timeTextDecode += 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]);
@@ -1523,11 +1550,14 @@ 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 });
 
-    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}`);
+    // 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}`);
 
     return { text, audioCodes };
   }
@@ -1568,14 +1598,11 @@ export class AudioModel {
       throw new Error('Vocoder not loaded - required for interleaved mode');
     }
 
-    // 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
+    // Timing accumulators
     let timePrefill = 0;
-    let timeLfmText = 0;
-    let timeLfmAudio = 0;
-    let timeDepthformer = 0;
+    let timeTextDecode = 0;
+    let timeAudioDecode = 0;
+    let timeVocoder = 0;
     let timeAudioEmbed = 0;
     let tStep;
 
@@ -1640,18 +1667,17 @@ export class AudioModel {
 
         tStep = performance.now();
         const frameCodes = await this.sampleAudioCodes(lastHidden, audioTemperature, audioTopK);
-        timeDepthformer += performance.now() - tStep;
+        timeVocoder += 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} (${audioCodes.length} frames collected)`);
+          log(`End of audio at step ${step}`);
           for (let i = 0; i < NUM_CODEBOOKS; i++) {
             frameCodes[i] = END_OF_AUDIO_TOKEN;
           }
@@ -1665,7 +1691,7 @@ export class AudioModel {
           }
 
           if (audioCodes.length % 50 === 0) {
-            log(`  Audio frames: ${audioCodes.length}`);
+            log(`Generated ${audioCodes.length} audio frames`);
           }
         }
 
@@ -1681,7 +1707,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));
-        timeLfmAudio += performance.now() - tStep;
+        timeAudioDecode += performance.now() - tStep;
         this.updateCache(this.cache, outputs);
 
       } else {
@@ -1697,19 +1723,18 @@ 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} (${textTokens.length} text tokens, ${audioCodes.length} audio frames)`);
+          log(`End of turn at step ${step}`);
           break;
         }
 
         // Check for <|text_end|> token
         if (token === SPECIAL_TOKENS.TEXT_END) {
-          log(`→ TEXT_END at step ${step}: audio-only phase begins`);
+          log(`Text end at step ${step}`);
           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;
         }
@@ -1727,7 +1752,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));
-        timeLfmText += performance.now() - tStep;
+        timeTextDecode += performance.now() - tStep;
         this.updateCache(this.cache, outputs);
       }
     }
@@ -1742,9 +1767,13 @@ export class AudioModel {
 
     const text = this.tokenizer.decode(textTokens, { skip_special_tokens: true });
 
-    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}`);
+    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}`);
 
     return { text, audioCodes };
   }
@@ -1874,8 +1903,7 @@ export class AudioModel {
     }
 
     const decodeStart = performance.now();
-
-    log(`Audio decode: ${audioCodes.length} frames → waveform`);
+    log(`Decoding ${audioCodes.length} audio frames...`);
 
     // ISTFT parameters (fixed for this model)
     const nFft = 1280;
@@ -1883,7 +1911,7 @@ export class AudioModel {
     const winLength = 1280;
     const nFftBins = nFft / 2 + 1;
 
-    // Transpose codes [T, 8] → [1, 8, T] for ONNX input
+    // Stack codes: [T, 8] -> [8, T] and add batch -> [1, 8, T]
     const T = audioCodes.length;
     const codesTransposed = new BigInt64Array(8 * T);
     for (let t = 0; t < T; t++) {
@@ -1892,18 +1920,18 @@ export class AudioModel {
       }
     }
 
-    // Run detokenizer ONNX: [1, 8, T] → [1, 6T, 1282]
+    // Run detokenizer: [1, 8, T] -> [1, T, 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;
-    const actualT = stftFeatures.dims[1];
-    const detokEnd = performance.now();
-    log(`  Detokenizer: [1,8,${T}] → [1,${actualT},1282] in ${(detokEnd - detokStart).toFixed(0)}ms`);
+    log(`Detokenizer: ${(performance.now() - detokStart).toFixed(0)}ms, STFT frames: ${stftFeatures.dims[1]}`);
 
-    // Split into magnitude + angle → complex spectrogram
+    // Get raw data - shape is [1, T, 1282], we need to skip batch dimension
     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);
@@ -1911,23 +1939,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 (inverse Short-Time Fourier Transform) → waveform
-    const pad = (winLength - hopLength) / 2;
+    // ISTFT with 'same' padding
     const istftStart = performance.now();
     const waveform = this.istftSamePadding(complexStft, nFft, hopLength, winLength, actualT);
-    const istftEnd = performance.now();
-    log(`  ISTFT: ${actualT} frames → ${waveform.length} samples in ${(istftEnd - istftStart).toFixed(0)}ms`);
+    log(`ISTFT: ${(performance.now() - istftStart).toFixed(0)}ms`);
 
-    // Find max/min
+    // Find max/min without spread operator (avoid stack overflow on large arrays)
     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)) {
@@ -1935,7 +1963,7 @@ export class AudioModel {
       return new Float32Array(0);
     }
 
-    // Normalize waveform to [-0.9, 0.9]
+    // Normalize to [-1, 1]
     let maxVal = Math.max(Math.abs(waveMax), Math.abs(waveMin));
     if (maxVal > 0) {
       for (let i = 0; i < waveform.length; i++) {
@@ -1945,8 +1973,7 @@ export class AudioModel {
       console.warn('ISTFT produced all-zero waveform');
     }
 
-    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`);
+    log(`Decoded audio: ${waveform.length} samples (${(waveform.length / 24000).toFixed(2)}s)`);
     return waveform;
   }