feat(stage-ui): better tts input chunking (#306)

* feat(stage-ui): better tts input chunking * chore: remove unused audio analyser * fix: words count * fix: use ctx.data instead --------- Co-authored-by: Neko <neko@ayaka.moe>
2025-07-21 14:43:56 +08:00
parent 4c2cade988
commit 42026d100b
7 changed files with 327 additions and 13 deletions
@@ -33,6 +33,7 @@ words:
  - cientos
  - cjkfonts
  - clippy
+  - clustr
  - collectblock
  - colorjs
  - Comfortaa
@@ -139,6 +139,7 @@
    "@types/three": "^0.178.1",
    "@unocss/reset": "^66.3.3",
    "@vitejs/plugin-vue": "^6.0.0",
+    "clustr": "^0.0.3",
    "histoire": "1.0.0-alpha.2",
    "unocss": "^66.3.3",
    "unplugin-yaml": "^3.0.1",
@@ -5,6 +5,8 @@ import { FieldCheckbox, FieldSelect } from '@proj-airi/ui'
 import { computed, onUnmounted, ref, watch } from 'vue'
 import { useI18n } from 'vue-i18n'

+import SpeechStreamingPlayground from './SpeechStreamingPlayground.vue'
+
 import { TestDummyMarker } from '../../Gadgets'

 const props = defineProps<{
@@ -204,6 +206,12 @@ defineExpose({
        {{ errorMessage }}
      </div>
      <audio v-if="audioUrl" ref="audioPlayer" :src="audioUrl" controls class="mt-2 w-full" />
+
+      <SpeechStreamingPlayground
+        :text="testText"
+        :voice="selectedVoice"
+        :generate-speech="generateSpeech"
+      />
    </div>
    <!-- Slot for additional provider-specific UI in the playground -->
    <slot />
@@ -0,0 +1,163 @@
+<script setup lang="ts">
+import type { TTSInputChunk } from '../../../utils/tts'
+
+import { animate } from 'animejs'
+import { ref } from 'vue'
+
+import { useQueue } from '../../../composables/queue'
+import { useMessageContentQueue } from '../../../composables/queues'
+import { useAudioContext } from '../../../stores/audio'
+import { chunkTTSInput } from '../../../utils/tts'
+
+const props = defineProps<{
+  text: string
+  // Provider-specific handlers (provided from parent)
+  generateSpeech: (input: string, voice: string, useSSML: boolean) => Promise<ArrayBuffer>
+  voice: string
+}>()
+
+const { audioContext } = useAudioContext()
+const nowSpeaking = ref(false)
+const ttsInputChunks = ref<TTSInputChunk[]>([])
+const speechGenerationIndex = ref(-1)
+
+const audioQueue = useQueue<{ audioBuffer: AudioBuffer, text: string }>({
+  handlers: [
+    (ctx) => {
+      return new Promise((resolve) => {
+        // Create an AudioBufferSourceNode
+        const source = audioContext.createBufferSource()
+        source.buffer = ctx.data.audioBuffer
+
+        // Connect the source to the AudioContext's destination (the speakers)
+        source.connect(audioContext.destination)
+
+        // Start playing the audio
+        nowSpeaking.value = true
+        source.start(0)
+        source.onended = () => {
+          nowSpeaking.value = false
+          resolve()
+        }
+      })
+    },
+  ],
+})
+
+async function handleSpeechGeneration(ctx: { data: string }) {
+  speechGenerationIndex.value++
+
+  try {
+    const input = ctx.data
+    const res = await props.generateSpeech(input, props.voice, false)
+
+    // Decode the ArrayBuffer into an AudioBuffer
+    const audioBuffer = await audioContext.decodeAudioData(res)
+    await audioQueue.add({ audioBuffer, text: ctx.data })
+  }
+  catch (error) {
+    console.error('Speech generation failed:', error)
+  }
+}
+
+const ttsQueue = useQueue<string>({
+  handlers: [
+    handleSpeechGeneration,
+  ],
+})
+
+const messageContentQueue = useMessageContentQueue(ttsQueue)
+
+async function testStreaming() {
+  await messageContentQueue.add(props.text)
+}
+
+async function testChunking() {
+  const chunks = []
+  for await (const chunk of chunkTTSInput(props.text, { boost: 1, minimumWords: 4, maximumWords: 12 })) {
+    chunks.push(chunk)
+  }
+  ttsInputChunks.value = chunks
+}
+</script>
+
+<template>
+  <div class="flex items-center gap-1 text-sm font-medium">
+    Streaming Playground
+  </div>
+  <div flex="~ row" gap-4>
+    <button
+      border="neutral-800 dark:neutral-200 solid 2" transition="border duration-250 ease-in-out"
+      rounded-lg px-4 text="neutral-100 dark:neutral-900" py-2 text-sm
+      bg="neutral-700 dark:neutral-300" @click="testChunking"
+    >
+      <div flex="~ row" items-center gap-2>
+        <div i-solar:round-double-alt-arrow-right-bold-duotone />
+        <span>Test chunking</span>
+      </div>
+    </button>
+
+    <button
+      v-if="ttsInputChunks.length > 0"
+      border="neutral-800 dark:neutral-200 solid 2" transition="border duration-250 ease-in-out"
+      rounded-lg px-4 text="neutral-100 dark:neutral-900" py-2 text-sm
+      bg="neutral-700 dark:neutral-300" @click="testStreaming"
+    >
+      <div flex="~ row" items-center gap-2>
+        <div i-solar:round-double-alt-arrow-right-bold-duotone />
+        <span>Test streaming</span>
+      </div>
+    </button>
+  </div>
+
+  <div flex="~ col gap-2 items-start" py-4>
+    <div
+      v-for="(chunk, i) in ttsInputChunks"
+      :key="i"
+      flex="~ row gap-2 items-center"
+    >
+      <div
+        flex="~ row gap-2 items-center"
+        rounded-xl px-2 py-1.5
+        :class="{
+          'bg-neutral-100 dark:bg-neutral-800': speechGenerationIndex < i,
+          'bg-neutral-200 dark:bg-neutral-700': speechGenerationIndex >= i,
+        }"
+      >
+        <span ml-1>{{ chunk.text }}</span>
+        <span
+          rounded-full px-2 py-.5 text-nowrap text-xs
+          b="~ dashed"
+          :class="{
+            'b-green text-green': chunk.reason === 'boost',
+            'b-orange text-orange': chunk.reason === 'limit',
+            'b-red text-red': chunk.reason === 'hard',
+            'b-purple text-purple': chunk.reason === 'flush',
+          }"
+        >
+          {{ chunk.words }} words,
+          {{ chunk.reason }}
+        </span>
+      </div>
+      <Transition
+        :css="false"
+        @enter="(el) => animate(el, {
+          opacity: [0, 1],
+          translateX: [10, 0],
+          duration: 200,
+          ease: 'inOut',
+        })"
+      >
+        <div
+          v-if="speechGenerationIndex >= i"
+          tag="div"
+          flex="~ row items-center gap-1"
+          text-sm
+        >
+          <div i-solar-check-circle-line-duotone />
+          <div>Queued</div>
+        </div>
+      </Transition>
+    </div>
+  </div>
+</template>
@@ -6,6 +6,7 @@ import { ref } from 'vue'

 import { llmInferenceEndToken } from '../constants'
 import { EMOTION_VALUES } from '../constants/emotions'
+import { chunkTTSInput } from '../utils/tts'
 import { useQueue } from './queue'

 export function useEmotionsMessageQueue(emotionsQueue: UseQueueReturn<Emotion>) {
@@ -116,19 +117,8 @@ export function useMessageContentQueue(ttsQueue: UseQueueReturn<string>) {
          return
        }

-        const endMarker = /[.?!]/
-        processed.value += ctx.data
-
-        while (processed.value) {
-          const endMarkerExecArray = endMarker.exec(processed.value)
-          if (!endMarkerExecArray || typeof endMarkerExecArray.index === 'undefined')
-            break
-
-          const before = processed.value.slice(0, endMarkerExecArray.index + 1)
-          const after = processed.value.slice(endMarkerExecArray.index + 1)
-
-          await ttsQueue.add(before)
-          processed.value = after
+        for await (const chunk of chunkTTSInput(ctx.data)) {
+          await ttsQueue.add(chunk.text)
        }
      },
    ],
@@ -0,0 +1,142 @@
+import type { ReaderLike } from 'clustr'
+
+import { readGraphemeClusters } from 'clustr'
+
+const keptPunctuations = new Set('?？!！')
+const hardPunctuations = new Set('.。?？!！…⋯～~「」\n\t\r')
+const softPunctuations = new Set(',，、–—:：;；《》')
+
+export interface TTSInputChunk {
+  text: string
+  words: number
+  reason: 'boost' | 'limit' | 'hard' | 'flush'
+}
+
+export interface TTSInputChunkOptions {
+  boost?: number
+  minimumWords?: number
+  maximumWords?: number
+}
+
+/**
+ * Processes the input string or UTF-8 byte stream reader into chunks suitable for TTS synthesis.
+ *
+ * @param input A string or a ReaderLike object that reads from an underlying UTF-8 byte stream.
+ * @param options
+ * @param options.boost Specifies the number of chunks to yield using greedier rules. This may help
+ *                      reduce the initial delay when processing long input text.
+ * @param options.minimumWords Minimum number of words in a chunk.
+ * @param options.maximumWords Maximum number of words in a chunk.
+ */
+export async function* chunkTTSInput(input: string | ReaderLike, options?: TTSInputChunkOptions): AsyncGenerator<TTSInputChunk, void, unknown> {
+  const {
+    boost = 2,
+    minimumWords = 4,
+    maximumWords = 12,
+  } = options ?? {}
+
+  const iterator = readGraphemeClusters(
+    typeof input === 'string'
+      ? new ReadableStream({
+          start(controller) {
+            controller.enqueue(new TextEncoder().encode(input))
+            controller.close()
+          },
+        }).getReader()
+      : input,
+  )
+
+  const segmenter = new Intl.Segmenter(undefined, { granularity: 'word' }) // I love Intl.Segmenter
+
+  let yieldCount = 0
+  let buffer = ''
+  let chunk = ''
+  let chunkWordsCount = 0
+
+  let previousValue: string | undefined
+  let current = await iterator.next()
+
+  while (!current.done) {
+    const value = current.value
+
+    if (value.length > 1) {
+      previousValue = value
+      current = await iterator.next()
+      continue
+    }
+
+    const hard = hardPunctuations.has(value)
+    const soft = softPunctuations.has(value)
+    const kept = keptPunctuations.has(value)
+
+    if (hard || soft) {
+      switch (value) {
+        case '.':
+        case ',': {
+          if (previousValue !== undefined && /\d/.test(previousValue)) {
+            const next = await iterator.next()
+            if (!next.done && next.value && /\d/.test(next.value)) {
+              // This dot could be a decimal point, so we skip it
+              previousValue = next.value
+              current = next
+              continue
+            }
+          }
+        }
+      }
+
+      if (buffer.length === 0) {
+        previousValue = value
+        current = await iterator.next()
+        continue
+      }
+
+      const words = [...segmenter.segment(buffer)].filter(w => w.isWordLike)
+
+      if (chunkWordsCount > minimumWords && chunkWordsCount + words.length > maximumWords) {
+        const text = kept ? chunk.trim() + value : chunk.trim()
+        yield {
+          text,
+          words: chunkWordsCount,
+          reason: 'limit',
+        }
+        yieldCount++
+        chunk = ''
+        chunkWordsCount = 0
+      }
+
+      chunk += buffer + value
+      chunkWordsCount += words.length
+      buffer = ''
+
+      if (hard || chunkWordsCount > maximumWords || yieldCount < boost) {
+        const text = chunk.trim()
+        yield {
+          text,
+          words: chunkWordsCount,
+          reason: hard ? 'hard' : chunkWordsCount > maximumWords ? 'limit' : 'boost',
+        }
+        yieldCount++
+        chunk = ''
+        chunkWordsCount = 0
+      }
+
+      previousValue = value
+      current = await iterator.next()
+      continue
+    }
+
+    buffer += value
+    previousValue = value
+    current = await iterator.next()
+  }
+
+  if (chunk.length > 0 || buffer.length > 0) {
+    const text = (chunk + buffer).trim()
+    yield {
+      text,
+      words: chunkWordsCount + [...segmenter.segment(buffer)].filter(w => w.isWordLike).length,
+      reason: 'flush',
+    }
+  }
+}
@@ -1419,6 +1419,9 @@ importers:
      '@vitejs/plugin-vue':
        specifier: ^6.0.0
        version: 6.0.0(vite@6.3.5(@types/node@24.0.13)(jiti@2.4.2)(less@4.3.0)(lightningcss@1.30.1)(terser@5.43.1)(tsx@4.20.3)(yaml@2.8.0))(vue@3.5.17(typescript@5.8.3))
+      clustr:
+        specifier: ^0.0.3
+        version: 0.0.3
      histoire:
        specifier: 1.0.0-alpha.2
        version: 1.0.0-alpha.2(@types/node@24.0.13)(bufferutil@4.0.9)(less@4.3.0)(lightningcss@1.30.1)(terser@5.43.1)(utf-8-validate@5.0.10)(vite@6.3.5(@types/node@24.0.13)(jiti@2.4.2)(less@4.3.0)(lightningcss@1.30.1)(terser@5.43.1)(tsx@4.20.3)(yaml@2.8.0))
@@ -7047,6 +7050,10 @@ packages:
    resolution: {integrity: sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==}
    engines: {node: '>=6'}

+  clustr@0.0.3:
+    resolution: {integrity: sha512-nGJZKUSi2K83migJ7sJj/JYNVbT0L81b7fFNcQ3XIgkXqkeG1tpO1NtISDjr3Unnjqsztzvp/p+8JgBhpAc87g==}
+    engines: {node: '>=22.10.0'}
+
  color-convert@2.0.1:
    resolution: {integrity: sha512-RRECPsj7iu/xb5oKYcsFHSppFNnsj/52OVTRKb4zP5onXwVF3zVmmToNcOfGC+CRDpfK/U584fMg38ZHCaElKQ==}
    engines: {node: '>=7.0.0'}
@@ -19426,6 +19433,8 @@ snapshots:
  clsx@2.1.1:
    optional: true

+  clustr@0.0.3: {}
+
  color-convert@2.0.1:
    dependencies:
      color-name: 1.1.4