From 1fa2aa3ff559dc42d89d00ab5eab8510067c6b04 Mon Sep 17 00:00:00 2001
From: Neko Ayaka <neko@ayaka.moe>
Date: Wed, 25 Dec 2024 17:22:39 +0800
Subject: [PATCH] chore: no top level await

---
 packages/moonshine-web/src/App.vue        |  43 ++--
 packages/moonshine-web/src/libs/types.ts  |  25 ++-
 packages/moonshine-web/src/libs/worker.ts | 250 +++++++++++++---------
 3 files changed, 194 insertions(+), 124 deletions(-)

diff --git a/packages/moonshine-web/src/App.vue b/packages/moonshine-web/src/App.vue
index 401271552..27118a9f6 100644
--- a/packages/moonshine-web/src/App.vue
+++ b/packages/moonshine-web/src/App.vue
@@ -1,5 +1,5 @@
 <script setup lang="ts">
-import type { MessageEvent, MessageEventInfo, MessageEventOutput, MessageEventStatus } from './libs/types'
+import type { MessageEvent, MessageEventBufferRequest, MessageEventInfo, MessageEventLoad, MessageEventOutput, MessageEventStatus } from './libs/types'
 import { TresCanvas } from '@tresjs/core'
 import { useWebWorker } from '@vueuse/core'
 import { ACESFilmicToneMapping, SRGBColorSpace } from 'three'
@@ -16,7 +16,7 @@ import { formatDate } from './utils'
 
 const status = ref<string | null>(null)
 const error = ref(null)
-const messages = ref<Array<MessageEventStatus | MessageEventInfo | MessageEventOutput>>([])
+const messages = ref<Array<MessageEventStatus | MessageEventInfo | MessageEventOutput | MessageEventBufferRequest | MessageEventLoad>>([])
 const frequency = ref(0)
 
 const { post, data } = useWebWorker<MessageEvent>(Worker, { type: 'module' })
@@ -31,13 +31,27 @@ watch(data, () => {
   }
   if (data.value.type === MessageType.Status) {
     status.value = data.value.message
-    messages.value = [...messages.value, data.value]
+    messages.value.push(data.value)
+
+    // pop out the other messages except the last status message
+    if (messages.value.length > 1) {
+      messages.value = messages.value.slice(-1)
+    }
   }
   else {
-    messages.value = [...messages.value, data.value]
+    messages.value.push(data.value)
+
+    // pop out the last message
+    if (messages.value.length > 1) {
+      messages.value = messages.value.slice(-1)
+    }
   }
 })
 
+onMounted(() => {
+  post({ type: MessageType.Load } satisfies MessageEventLoad)
+})
+
 onMounted(() => {
   // https://react.dev/learn/synchronizing-with-effects#fetching-data
   let ignore = false // Flag to track if the effect is active
@@ -107,7 +121,7 @@ onMounted(() => {
         const { buffer } = event.data
 
         // Dispatch buffer for voice activity detection
-        post({ buffer })
+        post({ type: MessageType.Request, buffer } satisfies MessageEventBufferRequest)
       }
     })
     .catch((err) => {
@@ -174,21 +188,20 @@ function downloadTranscript() {
     </template>
     <template v-else>
       <div class="absolute bottom-0 z-10 w-full overflow-hidden pb-8 text-center text-white">
-        <template v-for="(message, index) of messages" :key="index">
+        <TransitionGroup name="fade-up" tag="div">
           <div
+            v-for="(message) of messages" :key="message.message || ''"
             :initial="{ opacity: 0, y: 25 }"
             :enter="{ opacity: 1, y: 0 }"
             :duration="200"
             class="mb-1"
             :class="[message.type === 'output' ? 'text-5xl' : 'text-2xl text-green-300 font-light']"
           >
-            <Transition name="fade-up">
-              <div v-if="message.duration === 'until_next' && index === messages.length - 1">
-                {{ message.message }}
-              </div>
-            </Transition>
+            <div>
+              {{ message.message }}
+            </div>
           </div>
-        </template>
+        </TransitionGroup>
       </div>
       <TresCanvas window-size :alpha="true" :antialias="true" power-preference="high-performance" :output-color-space="SRGBColorSpace" :tone-mapping="ACESFilmicToneMapping">
         <TresPerspectiveCamera :position="[0, 0, 8]" :fov="75" :near="0.1" :far="1000" />
@@ -255,7 +268,11 @@ function downloadTranscript() {
   transition: all 0.5s ease-in-out;
 }
 
-.fade-up-enter-from,
+.fade-up-enter-from {
+  opacity: 0;
+  transform: translateY(25px);
+}
+
 .fade-up-leave-to {
   opacity: 0;
   transform: translateY(-25px);
diff --git a/packages/moonshine-web/src/libs/types.ts b/packages/moonshine-web/src/libs/types.ts
index a3f81357e..ebafabdef 100644
--- a/packages/moonshine-web/src/libs/types.ts
+++ b/packages/moonshine-web/src/libs/types.ts
@@ -2,11 +2,15 @@ export enum MessageType {
   Status = 'status',
   Output = 'output',
   Info = 'info',
+  Request = 'request',
+  Error = 'error',
+  Load = 'load',
 }
 
 export enum MessageStatus {
   RecordingStart = 'recording_start',
   RecordingEnd = 'recording_end',
+  Ready = 'ready',
 }
 
 export enum Duration {
@@ -17,7 +21,7 @@ export interface MessageEventStatus {
   type: MessageType.Status
   status: MessageStatus
   message: string
-  duration: Duration
+  duration?: Duration
 }
 
 export interface MessageEventOutput {
@@ -35,8 +39,21 @@ export interface MessageEventInfo {
   duration?: Duration.UntilNext
 }
 
-export interface MessageEventError {
-  error: unknown
+export interface MessageEventBufferRequest {
+  type: MessageType.Request
+  buffer: Float32Array<any>
+  message?: string
 }
 
-export type MessageEvent = MessageEventError | MessageEventStatus | MessageEventOutput | MessageEventInfo
+export interface MessageEventError {
+  type: MessageType.Error
+  error: unknown
+  message?: string
+}
+
+export interface MessageEventLoad {
+  type: MessageType.Load
+  message?: string
+}
+
+export type MessageEvent = MessageEventError | MessageEventStatus | MessageEventOutput | MessageEventInfo | MessageEventBufferRequest | MessageEventLoad
diff --git a/packages/moonshine-web/src/libs/worker.ts b/packages/moonshine-web/src/libs/worker.ts
index 361028ed5..6b5f7c30d 100644
--- a/packages/moonshine-web/src/libs/worker.ts
+++ b/packages/moonshine-web/src/libs/worker.ts
@@ -1,6 +1,6 @@
-/* eslint-disable antfu/no-top-level-await */
 /* eslint-disable no-restricted-globals */
-import type { MessageEventError, MessageEventInfo, MessageEventOutput, MessageEventStatus } from './types'
+import type { AutomaticSpeechRecognitionPipeline, PreTrainedModel } from '@huggingface/transformers'
+import type { MessageEvent as InternalMessageEvent, MessageEventBufferRequest, MessageEventError, MessageEventInfo, MessageEventOutput, MessageEventStatus } from './types'
 import { AutoModel, pipeline, Tensor } from '@huggingface/transformers'
 
 import {
@@ -16,51 +16,13 @@ import {
 import { supportsWebGPU } from '../utils'
 import { Duration, MessageStatus, MessageType } from './types'
 
-const device = (await supportsWebGPU()) ? 'webgpu' : 'wasm'
-self.postMessage({ type: MessageType.Info, message: `Using device: "${device}"` } satisfies MessageEventInfo)
-self.postMessage({
-  type: MessageType.Info,
-  message: 'Loading models...',
-  duration: Duration.UntilNext,
-} satisfies MessageEventInfo)
+export type DType = Record<string, Exclude<NonNullable<Required<Parameters<typeof pipeline>>[2]['dtype']>, string>[string]>
+export type Device = Extract<Exclude<NonNullable<Required<Parameters<typeof pipeline>>[2]['device']>, Record<string, any>>, 'webgpu' | 'wasm'>
+export type PretrainedConfig = NonNullable<Parameters<typeof AutoModel.from_pretrained>[1]>['config']
 
 // Load models
-const silero_vad = await AutoModel.from_pretrained(
-  'onnx-community/silero-vad',
-  {
-    config: { model_type: 'custom' },
-    dtype: 'fp32', // Full-precision
-  },
-).catch((error) => {
-  self.postMessage({ error } satisfies MessageEventError)
-  throw error
-})
-
-const DEVICE_DTYPE_CONFIGS = {
-  webgpu: {
-    encoder_model: 'fp32',
-    decoder_model_merged: 'q4',
-  },
-  wasm: {
-    encoder_model: 'fp32',
-    decoder_model_merged: 'q8',
-  },
-}
-
-const transcriber = await pipeline(
-  'automatic-speech-recognition',
-  'onnx-community/moonshine-base-ONNX', // or "onnx-community/whisper-tiny.en",
-  {
-    device,
-    dtype: DEVICE_DTYPE_CONFIGS[device],
-  },
-).catch((error) => {
-  self.postMessage({ error } satisfies MessageEventError)
-  throw error
-})
-
-await transcriber(new Float32Array(SAMPLE_RATE)) // Compile shaders
-self.postMessage({ type: 'status', status: 'ready', message: 'Ready!' })
+let silero_vad: PreTrainedModel
+let transcriber: AutomaticSpeechRecognitionPipeline
 
 // Transformers.js currently doesn't support simultaneous inference,
 // so we need to chain the inference promises.
@@ -77,19 +39,62 @@ let state = new Tensor('float32', new Float32Array(2 * 1 * 128), [2, 1, 128])
 // Whether we are in the process of adding audio to the buffer
 let isRecording = false
 
+// Track the number of samples after the last speech chunk
+let postSpeechSamples = 0
+
+const DEVICE_DTYPE_CONFIGS: Record<Device, DType> = {
+  webgpu: {
+    encoder_model: 'fp32',
+    decoder_model_merged: 'q4',
+  },
+  wasm: {
+    encoder_model: 'fp32',
+    decoder_model_merged: 'q8',
+  },
+}
+
+async function newVADModel() {
+  // Load models
+  return await AutoModel.from_pretrained(
+    'onnx-community/silero-vad',
+    {
+      config: { model_type: 'custom' } as PretrainedConfig,
+      dtype: 'fp32', // Full-precision
+    },
+  ).catch((error) => {
+    self.postMessage({ type: MessageType.Error, error } satisfies MessageEventError)
+    throw error
+  })
+}
+
+async function newAutomaticSpeechRecognitionPipeline(device: Device) {
+  return await pipeline(
+    'automatic-speech-recognition',
+    'onnx-community/moonshine-base-ONNX', // or "onnx-community/whisper-tiny.en",
+    {
+      device,
+      dtype: DEVICE_DTYPE_CONFIGS[device],
+    },
+  ).catch((error) => {
+    self.postMessage({ type: MessageType.Error, error } satisfies MessageEventError)
+    throw error
+  })
+}
+
 /**
  * Perform Voice Activity Detection (VAD)
  * @param {Float32Array} buffer The new audio buffer
  * @returns {Promise<boolean>} `true` if the buffer is speech, `false` otherwise.
  */
 async function vad(buffer: Float32Array<ArrayBuffer>) {
+  if (silero_vad === undefined) {
+    console.warn('VAD model not loaded yet')
+    return false
+  }
+
   const input = new Tensor('float32', buffer, [1, buffer.length])
-
-  const { stateN, output } = await (inferenceChain = inferenceChain.then(_ =>
-    silero_vad({ input, sr, state }),
-  ))
+  const { stateN, output } = await (inferenceChain = inferenceChain.then(_ => silero_vad({ input, sr, state })))
   state = stateN // Update state
-
   const isSpeech = output.data[0]
 
   // Use heuristics to determine if the buffer is speech or not
@@ -110,14 +115,16 @@ async function vad(buffer: Float32Array<ArrayBuffer>) {
  * @param {number} data.duration The duration of the speech segment
  */
 async function transcribe(buffer: Float32Array<any>, data: { start: number, end: number, duration: number }) {
-  const { text } = await (inferenceChain = inferenceChain.then(_ =>
-    transcriber(buffer),
-  ))
+  if (transcriber === undefined) {
+    console.warn('Transcriber model not loaded yet')
+    return
+  }
+
+  // @ts-expect-error - chain
+  const { text }: { text: string } = await (inferenceChain = inferenceChain.then(_ => transcriber(buffer)))
   self.postMessage({ type: MessageType.Output, buffer, message: text, ...data } satisfies MessageEventOutput)
 }
 
-// Track the number of samples after the last speech chunk
-let postSpeechSamples = 0
 function reset(offset = 0) {
   self.postMessage({
     type: MessageType.Status,
@@ -125,6 +132,7 @@ function reset(offset = 0) {
     message: 'Transcribing...',
     duration: Duration.UntilNext,
   } satisfies MessageEventStatus)
+
   BUFFER.fill(0, offset)
   bufferPointer = offset
   isRecording = false
@@ -134,10 +142,10 @@ function reset(offset = 0) {
 const prevBuffers: Array<Float32Array<ArrayBuffer>> = []
 
 function dispatchForTranscriptionAndResetAudioBuffer(overflow?: Float32Array<ArrayBuffer>) {
-  // Get start and end time of the speech segment, minus the padding
+// Get start and end time of the speech segment, minus the padding
   const now = Date.now()
   const end
-      = now - ((postSpeechSamples + SPEECH_PAD_SAMPLES) / SAMPLE_RATE) * 1000
+    = now - ((postSpeechSamples + SPEECH_PAD_SAMPLES) / SAMPLE_RATE) * 1000
   const start = end - (bufferPointer / SAMPLE_RATE) * 1000
   const duration = end - start
   const overflowLength = overflow?.length ?? 0
@@ -152,6 +160,7 @@ function dispatchForTranscriptionAndResetAudioBuffer(overflow?: Float32Array<Arr
     paddedBuffer.set(prev, offset)
     offset += prev.length
   }
+
   paddedBuffer.set(buffer, offset)
   transcribe(paddedBuffer, { start, end, duration })
 
@@ -163,76 +172,103 @@ function dispatchForTranscriptionAndResetAudioBuffer(overflow?: Float32Array<Arr
   reset(overflowLength)
 }
 
-self.onmessage = async (event) => {
-  const { buffer } = event.data as { buffer: Float32Array<ArrayBuffer> }
+async function load() {
+  const device = (await supportsWebGPU()) ? 'webgpu' : 'wasm'
+  self.postMessage({ type: MessageType.Info, message: `Using device: "${device}"` } satisfies MessageEventInfo)
+  self.postMessage({
+    type: MessageType.Info,
+    message: 'Loading models...',
+    duration: Duration.UntilNext,
+  } satisfies MessageEventInfo)
 
-  const wasRecording = isRecording // Save current state
-  const isSpeech = await vad(buffer)
+  // Load models
+  silero_vad = await newVADModel()
+  transcriber = await newAutomaticSpeechRecognitionPipeline(device)
 
-  if (!wasRecording && !isSpeech) {
+  await transcriber(new Float32Array(SAMPLE_RATE)) // Compile shaders
+  self.postMessage({ type: 'status', status: 'ready', message: 'Ready!' })
+
+  self.onmessage = async (event) => {
+    const { buffer } = event.data as MessageEventBufferRequest
+
+    const wasRecording = isRecording // Save current state
+    const isSpeech = await vad(buffer)
+
+    if (!wasRecording && !isSpeech) {
     // We are not recording, and the buffer is not speech,
     // so we will probably discard the buffer. So, we insert
     // into a FIFO queue with maximum size of PREV_BUFFER_SIZE
-    if (prevBuffers.length >= MAX_NUM_PREV_BUFFERS) {
+      if (prevBuffers.length >= MAX_NUM_PREV_BUFFERS) {
       // If the queue is full, we discard the oldest buffer
-      prevBuffers.shift()
+        prevBuffers.shift()
+      }
+
+      prevBuffers.push(buffer)
+      return
     }
 
-    prevBuffers.push(buffer)
-    return
-  }
-
-  const remaining = BUFFER.length - bufferPointer
-  if (buffer.length >= remaining) {
+    const remaining = BUFFER.length - bufferPointer
+    if (buffer.length >= remaining) {
     // The buffer is larger than (or equal to) the remaining space in the global buffer,
     // so we perform transcription and copy the overflow to the global buffer
-    BUFFER.set(buffer.subarray(0, remaining), bufferPointer)
-    bufferPointer += remaining
+      BUFFER.set(buffer.subarray(0, remaining), bufferPointer)
+      bufferPointer += remaining
 
-    // Dispatch the audio buffer
-    const overflow = buffer.subarray(remaining)
-    dispatchForTranscriptionAndResetAudioBuffer(overflow)
-    return
-  }
-  else {
+      // Dispatch the audio buffer
+      const overflow = buffer.subarray(remaining)
+      dispatchForTranscriptionAndResetAudioBuffer(overflow)
+      return
+    }
+    else {
     // The buffer is smaller than the remaining space in the global buffer,
     // so we copy it to the global buffer
-    BUFFER.set(buffer, bufferPointer)
-    bufferPointer += buffer.length
-  }
-
-  if (isSpeech) {
-    if (!isRecording) {
-      // Indicate start of recording
-      self.postMessage({
-        type: MessageType.Status,
-        status: MessageStatus.RecordingStart,
-        message: 'Listening...',
-        duration: Duration.UntilNext,
-      } satisfies MessageEventStatus)
+      BUFFER.set(buffer, bufferPointer)
+      bufferPointer += buffer.length
     }
-    // Start or continue recording
-    isRecording = true
-    postSpeechSamples = 0 // Reset the post-speech samples
-    return
-  }
 
-  postSpeechSamples += buffer.length
+    if (isSpeech) {
+      if (!isRecording) {
+      // Indicate start of recording
+        self.postMessage({
+          type: MessageType.Status,
+          status: MessageStatus.RecordingStart,
+          message: 'Listening...',
+          duration: Duration.UntilNext,
+        } satisfies MessageEventStatus)
+      }
+      // Start or continue recording
+      isRecording = true
+      postSpeechSamples = 0 // Reset the post-speech samples
+      return
+    }
 
-  // At this point we're confident that we were recording (wasRecording === true), but the latest buffer is not speech.
-  // So, we check whether we have reached the end of the current audio chunk.
-  if (postSpeechSamples < MIN_SILENCE_DURATION_SAMPLES) {
+    postSpeechSamples += buffer.length
+
+    // At this point we're confident that we were recording (wasRecording === true), but the latest buffer is not speech.
+    // So, we check whether we have reached the end of the current audio chunk.
+    if (postSpeechSamples < MIN_SILENCE_DURATION_SAMPLES) {
     // There was a short pause, but not long enough to consider the end of a speech chunk
     // (e.g., the speaker took a breath), so we continue recording
-    return
-  }
+      return
+    }
 
-  if (bufferPointer < MIN_SPEECH_DURATION_SAMPLES) {
+    if (bufferPointer < MIN_SPEECH_DURATION_SAMPLES) {
     // The entire buffer (including the new chunk) is smaller than the minimum
     // duration of a speech chunk, so we can safely discard the buffer.
-    reset()
-    return
-  }
+      reset()
+      return
+    }
 
-  dispatchForTranscriptionAndResetAudioBuffer()
+    dispatchForTranscriptionAndResetAudioBuffer()
+  }
 }
+
+self.addEventListener('message', (event) => {
+  const { type } = event.data as InternalMessageEvent
+
+  switch (type) {
+    case MessageType.Load:
+      load()
+      break
+  }
+})