correct the render

.gitignore

@@ -0,0 +1 @@
+node_modules/

audio-model.js

@@ -1289,6 +1289,7 @@ export class AudioModel {
 
     logReset();
     log('=== Interleaved Generation ===');
+    log('Cache state:', this.cache ? `exists (seq_len=${this.cacheSeqLen})` : 'null (new conversation)');
     log('Audio samples:', audioData.length, 'Sample rate:', sampleRate);
 
     if (!this.audioEncoderSession) {
@@ -1544,6 +1545,115 @@ export class AudioModel {
     return { text, audioCodes };
   }
 
+  /**
+   * Generate text-only response (for follow-up turns without audio).
+   * Uses the stateful KV cache from previous interleaved turns.
+   *
+   * @param {string} userText - User's text input
+   * @param {object} options - Generation options
+   * @returns {object} - { text: string }
+   */
+  async generateTextOnly(userText, options = {}) {
+    const {
+      maxNewTokens = 256,
+      temperature = 0.7,
+      systemPrompt = 'You are a helpful assistant.',
+      onToken,
+    } = options;
+
+    logReset();
+    log('=== Text-Only Generation ===');
+    log('Cache state:', this.cache ? `exists (seq_len=${this.cacheSeqLen})` : 'null (new conversation)');
+    log('User text:', userText);
+
+    if (!this.embedTokensWeight) {
+      throw new Error('embed_tokens not loaded');
+    }
+
+    const { hiddenSize } = this.embedTokensWeight;
+
+    // Build prompt based on whether we have existing cache
+    let inputEmbeds;
+    let newSeqLen;
+
+    if (this.cache === null) {
+      // First turn: include system message
+      log('First turn - initializing conversation');
+      this.cache = this.initializeCache();
+      this.cacheSeqLen = 0;
+
+      const promptText = `<|startoftext|><|im_start|>system\n${systemPrompt}<|im_end|>\n<|im_start|>user\n${userText}<|im_end|>\n<|im_start|>assistant\n`;
+      const promptIds = Array.from(this.tokenizer.encode(promptText, { add_special_tokens: false }));
+      inputEmbeds = this.getTextEmbeddings(promptIds);
+      newSeqLen = promptIds.length;
+    } else {
+      // Continuation: just user turn
+      log(`Continuing conversation (cache seq_len=${this.cacheSeqLen})`);
+
+      const turnText = `<|im_start|>user\n${userText}<|im_end|>\n<|im_start|>assistant\n`;
+      const turnIds = Array.from(this.tokenizer.encode(turnText, { add_special_tokens: false }));
+      inputEmbeds = this.getTextEmbeddings(turnIds);
+      newSeqLen = turnIds.length;
+    }
+
+    // Run prefill
+    const totalLen = this.cacheSeqLen + newSeqLen;
+    const attentionMask = new ort.Tensor('int64', new BigInt64Array(totalLen).fill(1n), [1, totalLen]);
+
+    let { logits, outputs } = await this.runDecoder(inputEmbeds, attentionMask, this.cache);
+    this.updateCache(this.cache, outputs);
+    this.cacheSeqLen = totalLen;
+
+    // Generate tokens
+    const textTokens = [];
+    let currentLen = totalLen;
+
+    for (let i = 0; i < maxNewTokens; i++) {
+      const logitsData = logits.data;
+      const seqLen = logits.dims[1];
+      const lastLogits = new Float32Array(this.vocabSize);
+      const offset = (seqLen - 1) * this.vocabSize;
+      for (let j = 0; j < this.vocabSize; j++) {
+        lastLogits[j] = logitsData[offset + j];
+      }
+      const nextToken = this.sampleToken(lastLogits, temperature);
+
+      // Check for stop tokens
+      if (nextToken === this.tokenizer.eos_token_id || nextToken === SPECIAL_TOKENS.IM_END) {
+        log('Stop token reached');
+        break;
+      }
+
+      textTokens.push(nextToken);
+
+      if (onToken) {
+        const text = this.tokenizer.decode(textTokens, { skip_special_tokens: true });
+        onToken(text, nextToken);
+      }
+
+      // Get embedding for next token
+      const nextEmbeds = this.getTextEmbeddings([nextToken]);
+      currentLen++;
+      const nextMask = new ort.Tensor('int64', new BigInt64Array(currentLen).fill(1n), [1, currentLen]);
+      ({ logits, outputs } = await this.runDecoder(nextEmbeds, nextMask, this.cache));
+      this.updateCache(this.cache, outputs);
+    }
+
+    // Feed <|im_end|> to close turn
+    const imEndEmbeds = this.getTextEmbeddings([SPECIAL_TOKENS.IM_END]);
+    currentLen++;
+    const finalMask = new ort.Tensor('int64', new BigInt64Array(currentLen).fill(1n), [1, currentLen]);
+    ({ outputs } = await this.runDecoder(imEndEmbeds, finalMask, this.cache));
+    this.updateCache(this.cache, outputs);
+    this.cacheSeqLen = currentLen;
+
+    const text = this.tokenizer.decode(textTokens, { skip_special_tokens: true });
+    log(`Generated ${textTokens.length} tokens: "${text}"`);
+    log(`Cache seq_len: ${this.cacheSeqLen}`);
+
+    return { text };
+  }
+
   /**
    * Decode audio codes to waveform using audio detokenizer + ISTFT
    * @param {number[][]} audioCodes - Array of [8] codebook values per frame

main.js

@@ -62,6 +62,24 @@ let audioChunks = [];
 // ============================================================================
 
 function createWavBlob(samples, sampleRate) {
+  // Debug: check waveform statistics
+  let min = Infinity, max = -Infinity, sum = 0, nonZero = 0;
+  for (let i = 0; i < samples.length; i++) {
+    const v = samples[i];
+    if (v < min) min = v;
+    if (v > max) max = v;
+    sum += Math.abs(v);
+    if (Math.abs(v) > 0.001) nonZero++;
+  }
+  console.log('WAV input stats:', {
+    length: samples.length,
+    min: min.toFixed(6),
+    max: max.toFixed(6),
+    avgAbs: (sum / samples.length).toFixed(6),
+    nonZeroSamples: nonZero,
+    percentNonZero: ((nonZero / samples.length) * 100).toFixed(1) + '%'
+  });
+
   const numChannels = 1;
   const bitsPerSample = 16;
   const bytesPerSample = bitsPerSample / 8;
@@ -104,6 +122,16 @@ function createWavBlob(samples, sampleRate) {
   return new Blob([buffer], { type: 'audio/wav' });
 }
 
+// Test function to verify WAV creation works
+function createTestToneBlob(durationSec = 1, frequency = 440, sampleRate = 24000) {
+  const numSamples = Math.floor(durationSec * sampleRate);
+  const samples = new Float32Array(numSamples);
+  for (let i = 0; i < numSamples; i++) {
+    samples[i] = 0.5 * Math.sin(2 * Math.PI * frequency * i / sampleRate);
+  }
+  return createWavBlob(samples, sampleRate);
+}
+
 // ============================================================================
 // UI Helpers
 // ============================================================================
@@ -441,22 +469,19 @@ async function generate(userMessage) {
         if (waveform.length > 0) {
           generatedText = result.textOutput || `Generated ${result.audioCodes.length} audio frames (${(waveform.length / 24000).toFixed(2)}s)`;
 
-          // Create audio player
-          const audioMsgEl = document.createElement('div');
-          audioMsgEl.className = 'message assistant';
+          // Create audio player inline with the message
           const wavBlob = createWavBlob(waveform, 24000);
-          console.log('TTS WAV blob created:', wavBlob.size, 'bytes');
+          console.log('TTS WAV blob created:', wavBlob.size, 'bytes, duration:', (waveform.length / 24000).toFixed(2), 's');
           const audioUrl = URL.createObjectURL(wavBlob);
 
-          const audioEl = document.createElement('audio');
-          audioEl.controls = true;
-          const sourceEl = document.createElement('source');
-          sourceEl.src = audioUrl;
-          sourceEl.type = 'audio/wav';
-          audioEl.appendChild(sourceEl);
-
-          audioMsgEl.appendChild(audioEl);
-          chatContainer.appendChild(audioMsgEl);
+          // Add audio element to the existing message
+          const audioContainer = document.createElement('div');
+          audioContainer.style.marginTop = '0.75rem';
+          audioContainer.innerHTML = `
+            <audio controls preload="auto" src="${audioUrl}" style="width:100%;max-width:360px;display:block;"></audio>
+            <a href="${audioUrl}" download="generated_audio.wav" style="display:block;font-size:0.7rem;margin-top:0.25rem;color:#666;">Download WAV (${(waveform.length / 24000).toFixed(1)}s)</a>
+          `;
+          msgEl.appendChild(audioContainer);
           chatContainer.scrollTop = chatContainer.scrollHeight;
         } else {
           generatedText = '[Audio decoding failed - no waveform generated]';
@@ -503,34 +528,54 @@ async function generate(userMessage) {
             generatedText = `Generated ${result.audioCodes.length} audio frames`;
           }
 
-          // Create audio player
-          const audioMsgEl = document.createElement('div');
-          audioMsgEl.className = 'message assistant';
+          // Create audio player inline with the message
           const wavBlob = createWavBlob(waveform, 24000);
-          console.log('WAV blob created:', wavBlob.size, 'bytes');
+          console.log('WAV blob created:', wavBlob.size, 'bytes, duration:', (waveform.length / 24000).toFixed(2), 's');
           const audioUrl = URL.createObjectURL(wavBlob);
 
-          const audioEl = document.createElement('audio');
-          audioEl.controls = true;
-          const sourceEl = document.createElement('source');
-          sourceEl.src = audioUrl;
-          sourceEl.type = 'audio/wav';
-          audioEl.appendChild(sourceEl);
-
-          audioMsgEl.appendChild(audioEl);
-          chatContainer.appendChild(audioMsgEl);
+          // Add audio element to the existing message
+          const audioContainer = document.createElement('div');
+          audioContainer.style.marginTop = '0.75rem';
+          audioContainer.innerHTML = `
+            <audio controls preload="auto" src="${audioUrl}" style="width:100%;max-width:360px;display:block;"></audio>
+            <a href="${audioUrl}" download="generated_audio.wav" style="display:block;font-size:0.7rem;margin-top:0.25rem;color:#666;">Download WAV (${(waveform.length / 24000).toFixed(1)}s)</a>
+          `;
+          msgEl.appendChild(audioContainer);
           chatContainer.scrollTop = chatContainer.scrollHeight;
         } else {
           console.warn('Waveform decoding returned empty result');
         }
       }
 
-    } else {
+    } else if (currentMode === 'interleaved' && userMessage) {
+      // Text-only follow-up in interleaved mode
       showSpinner('Generating response...');
-      generatedText = await audioModel.generate(messages, {
+      const result = await audioModel.generateTextOnly(userMessage, {
         maxNewTokens: 256,
-        onToken: onTokenCallback,
+        onToken: (text, tokenId) => {
+          generatedText = text;
+          tokenCount = text.length;
+          textEl.textContent = text;
+          chatContainer.scrollTop = chatContainer.scrollHeight;
+          const elapsed = ((performance.now() - startTime) / 1000).toFixed(1);
+          updateSpinner('Generating...', `${tokenCount} chars · ${elapsed}s`);
+        },
       });
+      generatedText = result.text || '';
+
+    } else if (userMessage) {
+      // Fallback text-only generation
+      showSpinner('Generating response...');
+      const result = await audioModel.generateTextOnly(userMessage, {
+        maxNewTokens: 256,
+        onToken: (text, tokenId) => {
+          generatedText = text;
+          tokenCount = text.length;
+          textEl.textContent = text;
+          chatContainer.scrollTop = chatContainer.scrollHeight;
+        },
+      });
+      generatedText = result.text || '';
     }
 
     generatedText = generatedText.replace(/<\|im_end\|>$/g, '').trim();