Add TTS and tokenizer

Author: nbsp

bun.lockb

@@ -24,6 +24,7 @@
     "typescript": "^5.0.0"
   },
   "dependencies": {
+    "commander": "^12.0.0",
     "@livekit/protocol": "^1.12.0"
   }
 }

package.json

@@ -2,6 +2,6 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-export { VAD, VADEventType, VADStream } from './vad';
-export { Plugin } from './plugin';
-export { version } from './version';
+export * from './vad';
+export * from './plugin';
+export * from './version';

src/index.ts

@@ -0,0 +1,22 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+export interface SegmentedSentence {
+  text: string;
+}
+
+export abstract class SentenceTokenizer {
+  abstract tokenize(text: string, language: string | undefined): SegmentedSentence[];
+  abstract stream(language: string | undefined): SentenceStream;
+}
+
+export abstract class SentenceStream implements IterableIterator<SegmentedSentence> {
+  abstract pushText(text: string): void;
+  abstract flush(): Promise<void>;
+  async close(): Promise<void> {}
+  abstract next(): IteratorResult<SegmentedSentence>;
+  [Symbol.iterator](): SentenceStream {
+    return this;
+  }
+}

src/tokenize.ts

@@ -0,0 +1,6 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+export { TTS, SynthesisEvent, SynthesisEventType, SynthesizedAudio, SynthesizeStream } from './tts';
+export { StreamAdapter, StreamAdapterWrapper } from './stream_adapter';

src/tts/index.ts

@@ -0,0 +1,91 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+import { TTS, SynthesisEvent, SynthesisEventType, SynthesizedAudio, SynthesizeStream } from './tts';
+import { SentenceStream, SentenceTokenizer } from '../tokenize';
+
+export class StreamAdapterWrapper extends SynthesizeStream {
+  closed: boolean;
+  tts: TTS;
+  sentenceStream: SentenceStream;
+  eventQueue: (SynthesisEvent | undefined)[];
+  task: {
+    run: Promise<void>;
+    cancel: () => void;
+  };
+
+  constructor(tts: TTS, sentenceStream: SentenceStream) {
+    super();
+    this.closed = false;
+    this.tts = tts;
+    this.sentenceStream = sentenceStream;
+    this.eventQueue = [];
+    this.task = {
+      run: new Promise((_, reject) => {
+        this.run(reject);
+      }),
+      cancel: () => {},
+    };
+  }
+
+  async run(reject: (arg: Error) => void) {
+    while (!this.closed) {
+      this.task.cancel = () => {
+        this.closed = true;
+        reject(new Error('cancelled'));
+      };
+      for await (const sentence of this.sentenceStream) {
+        const audio = await this.tts.synthesize(sentence.text);
+        this.eventQueue.push(new SynthesisEvent(SynthesisEventType.STARTED));
+        this.eventQueue.push(new SynthesisEvent(SynthesisEventType.AUDIO, audio));
+        this.eventQueue.push(new SynthesisEvent(SynthesisEventType.FINISHED));
+      }
+    }
+  }
+
+  pushText(token: string) {
+    this.sentenceStream.pushText(token);
+  }
+
+  async flush() {
+    await this.sentenceStream.flush();
+  }
+
+  next(): IteratorResult<SynthesisEvent> {
+    const event = this.eventQueue.shift();
+    if (event) {
+      return { done: false, value: event };
+    } else {
+      return { done: true, value: undefined };
+    }
+  }
+
+  async close(): Promise<void> {
+    this.task.cancel();
+    try {
+      await this.task.run;
+    } finally {
+      this.eventQueue.push(undefined);
+    }
+  }
+}
+
+export class StreamAdapter extends TTS {
+  tts: TTS;
+  tokenizer: SentenceTokenizer;
+
+  constructor(tts: TTS, tokenizer: SentenceTokenizer) {
+    super(true);
+    this.tts = tts;
+    this.tokenizer = tokenizer;
+  }
+
+  synthesize(text: string): Promise<SynthesizedAudio> {
+    return this.tts.synthesize(text);
+  }
+
+  stream() {
+    return new StreamAdapterWrapper(this.tts, this.tokenizer.stream(undefined));
+  }
+}

src/tts/stream_adapter.ts

@@ -0,0 +1,57 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+import { AudioFrame } from '@livekit/rtc-node';
+
+export interface SynthesizedAudio {
+  text: string;
+  data: AudioFrame;
+}
+
+export enum SynthesisEventType {
+  STARTED = 0,
+  AUDIO = 1,
+  FINISHED = 2,
+}
+
+export class SynthesisEvent {
+  type: SynthesisEventType;
+  audio: SynthesizedAudio | undefined;
+
+  constructor(type: SynthesisEventType, audio: SynthesizedAudio | undefined = undefined) {
+    this.type = type;
+    this.audio = audio;
+  }
+}
+
+export abstract class SynthesizeStream implements IterableIterator<SynthesisEvent> {
+  abstract pushText(token: string | undefined): void;
+
+  markSegmentEnd() {
+    this.pushText(undefined);
+  }
+
+  abstract close(wait: boolean): Promise<void>;
+  abstract next(): IteratorResult<SynthesisEvent>;
+
+  [Symbol.iterator](): SynthesizeStream {
+    return this;
+  }
+}
+
+export abstract class TTS {
+  #streamingSupported: boolean;
+
+  constructor(streamingSupported: boolean) {
+    this.#streamingSupported = streamingSupported;
+  }
+
+  abstract synthesize(text: string): Promise<SynthesizedAudio>;
+
+  abstract stream(): SynthesizeStream;
+
+  get streamingSupported(): boolean {
+    return this.#streamingSupported;
+  }
+}

src/tts/tts.ts

@@ -10,7 +10,7 @@ export enum VADEventType {
   END_OF_SPEECH = 3,
 }
 
-interface VADEvent {
+export interface VADEvent {
   type: VADEventType;
   samplesIndex: number;
   duration: number;
@@ -33,11 +33,11 @@ export abstract class VAD {
   }): VADStream;
 }
 
-export abstract class VADStream {
+export abstract class VADStream implements IterableIterator<VADEvent> {
   abstract pushFrame(frame: AudioFrame): void;
-  abstract aclose(wait: boolean): Promise<void>;
-  abstract anext(): Promise<VADEvent>;
-  private aiter(): VADStream {
+  abstract close(wait: boolean): Promise<void>;
+  abstract next(): IteratorResult<VADEvent>;
+  [Symbol.iterator](): VADStream {
     return this;
   }
 }