webapp/src/worker.ts

import { APP_CONFIG } from './config'
import { configureOrt, createTtsSessions, ensureWebGpuSupport, fetchCondEmb, fetchTokenizerJson } from './onnx'
import { runFinnishPipeline, type TtsSessions } from './pipeline'
import { createTokenizer, type BrowserTokenizer } from './tokenizer'
import type { WorkerRequest, WorkerResponse } from './types'

let sessionsPromise: Promise<TtsSessions> | null = null
let tokenizerPromise: Promise<BrowserTokenizer> | null = null
let condEmbPromise: Promise<Float32Array> | null = null
const workerScope = self as typeof self & {
  postMessage: (message: WorkerResponse, transfer?: Transferable[]) => void
  onmessage: ((event: MessageEvent<WorkerRequest>) => void | Promise<void>) | null
}

function postMessageToMain(message: WorkerResponse, transfer: Transferable[] = []): void {
  workerScope.postMessage(message, transfer)
}

function reportStatus(message: string): void {
  postMessageToMain({ type: 'status', message })
}

async function ensureInitialized(): Promise<{
  sessions: TtsSessions
  tokenizer: BrowserTokenizer
  condEmb: Float32Array
}> {
  ensureWebGpuSupport()
  configureOrt()

  sessionsPromise ??= createTtsSessions(reportStatus)
  tokenizerPromise ??= fetchTokenizerJson().then((json) => createTokenizer(json))
  condEmbPromise ??= fetchCondEmb()

  const [sessions, tokenizer, condEmb] = await Promise.all([
    sessionsPromise,
    tokenizerPromise,
    condEmbPromise,
  ])
  postMessageToMain({ type: 'ready' })
  return { sessions, tokenizer, condEmb }
}

async function handleSpeakRequest(message: Extract<WorkerRequest, { type: 'speak' }>): Promise<void> {
  const { sessions, tokenizer, condEmb } = await ensureInitialized()
  const result = await runFinnishPipeline({
    sessions,
    tokenizer,
    condEmb,
    referenceAudio: message.referenceAudio,
    text: message.text,
    reportStatus,
  })

  postMessageToMain(
    {
      type: 'result',
      audio: result.audio,
      sampleRate: APP_CONFIG.sampleRate,
      speechTokenCount: result.speechTokenCount,
    },
    [result.audio.buffer],
  )
}

workerScope.onmessage = async (event: MessageEvent<WorkerRequest>) => {
  try {
    if (event.data.type === 'speak') {
      await handleSpeakRequest(event.data)
    }
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error)
    postMessageToMain({ type: 'error', message })
  }
}