diff --git a/requirements.txt b/requirements.txt index 6188992d0..04c6a29a7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +autopep8==2.0.4 build==1.0.3 packaging==23.2 -pyproject_hooks==1.0.0 +pyproject_hooks==1.0.0 \ No newline at end of file diff --git a/src/dailyai/requirements.txt b/src/dailyai/requirements.txt index 53d28d6fd..66ffbd0bb 100644 --- a/src/dailyai/requirements.txt +++ b/src/dailyai/requirements.txt @@ -1,2 +1,3 @@ Pillow==10.1.0 -typing_extensions==4.9.0 \ No newline at end of file +typing_extensions==4.9.0 +faster-whisper==0.10.0 \ No newline at end of file diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py index 684912989..24c676709 100644 --- a/src/dailyai/services/ai_services.py +++ b/src/dailyai/services/ai_services.py @@ -1,5 +1,9 @@ +import array import asyncio +import io import logging +import math +import wave from dailyai.queue_frame import ( AudioQueueFrame, @@ -11,7 +15,7 @@ from dailyai.queue_frame import ( ) from abc import abstractmethod -from typing import AsyncGenerator, AsyncIterable, Iterable +from typing import AsyncGenerator, AsyncIterable, BinaryIO, Iterable from dataclasses import dataclass @@ -150,9 +154,40 @@ class ImageGenService(AIService): (url, image_data) = await self.run_image_gen(frame.text) yield ImageQueueFrame(url, image_data) +class STTService(AIService): + """STTService is a base class for speech-to-text services.""" + + _frame_rate: int + def __init__(self, frame_rate: int = 16000, **kwargs): + super().__init__(**kwargs) + self._frame_rate = frame_rate + + + @abstractmethod + async def run_stt(self, audio: BinaryIO) -> str: + """Returns transcript as a string""" + pass + + async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]: + """Processes a frame of audio data, either buffering or transcribing it.""" + if not isinstance(frame, AudioQueueFrame): + return + + data = frame.data + content = io.BufferedRandom(io.BytesIO()) + ww = wave.open(self._content, "wb") + ww.setnchannels(1) + ww.setsampwidth(2) + ww.setframerate(self._frame_rate) + ww.writeframesraw(data) + ww.close() + content.seek(0) + text = await self.run_stt(content) + yield TextQueueFrame(text) @dataclass class AIServiceConfig: tts: TTSService image: ImageGenService llm: LLMService + stt: STTService diff --git a/src/dailyai/services/daily_transport_service.py b/src/dailyai/services/daily_transport_service.py index f1eac2bbd..aed5d9c9e 100644 --- a/src/dailyai/services/daily_transport_service.py +++ b/src/dailyai/services/daily_transport_service.py @@ -31,6 +31,10 @@ from daily import ( class DailyTransportService(EventHandler): _daily_initialized = False _lock = threading.Lock() + + speaker_enabled: bool + speaker_sample_rate: int + def __init__( self, room_url: str, @@ -38,6 +42,9 @@ class DailyTransportService(EventHandler): bot_name: str, duration: float = 10, min_others_count: int = 1, + start_transcription: bool = True, + speaker_enabled: bool = False, + speaker_sample_rate: int = 16000, ): super().__init__() self.bot_name: str = bot_name @@ -46,6 +53,7 @@ class DailyTransportService(EventHandler): self.duration: float = duration self.expiration = time.time() + duration * 60 self.min_others_count = min_others_count + self.start_transcription = start_transcription # This queue is used to marshal frames from the async send queue to the thread that emits audio & video. # We need this to maintain the asynchronous behavior of asyncio queues -- to give async functions @@ -61,6 +69,8 @@ class DailyTransportService(EventHandler): self.camera_width = 1024 self.camera_height = 768 self.camera_enabled = False + self.speaker_enabled = speaker_enabled + self.speaker_sample_rate = speaker_sample_rate self.send_queue = asyncio.Queue() self.receive_queue = asyncio.Queue() @@ -144,9 +154,11 @@ class DailyTransportService(EventHandler): "camera", width=self.camera_width, height=self.camera_height, color_format="RGB" ) - self.speaker: VirtualSpeakerDevice = Daily.create_speaker_device( - "speaker", sample_rate=16000, channels=1 - ) + if self.speaker_enabled: + self.speaker: VirtualSpeakerDevice = Daily.create_speaker_device( + "speaker", sample_rate=self.speaker_sample_rate, channels=1 + ) + Daily.select_speaker_device("speaker") self.image: bytes | None = None self.camera_thread = Thread(target=self.run_camera, daemon=True) @@ -156,8 +168,6 @@ class DailyTransportService(EventHandler): self.frame_consumer_thread = Thread(target=self.frame_consumer, daemon=True) self.frame_consumer_thread.start() - Daily.select_speaker_device("speaker") - self.client.set_user_name(self.bot_name) self.client.join(self.room_url, self.token, completion=self.call_joined) self.my_participant_id = self.client.participants()["local"]["id"] @@ -201,9 +211,20 @@ class DailyTransportService(EventHandler): } ) - if self.token: + if self.token and self.start_transcription: self.client.start_transcription(self.transcription_settings) + + def _receive_audio(self): + """Receive audio from the Daily call and put it on the receive queue""" + seconds = 1 + desired_frame_count = self.speaker_sample_rate * seconds + while True: + buffer = self.speaker.read_frames(desired_frame_count) + if len(buffer) > 0: + frame = AudioQueueFrame(buffer) + asyncio.run_coroutine_threadsafe(self.receive_queue.put(frame), self.loop) + async def get_receive_frames(self): while True: frame = await self.receive_queue.get() @@ -266,6 +287,9 @@ class DailyTransportService(EventHandler): def call_joined(self, join_data, client_error): self.logger.info(f"Call_joined: {join_data}, {client_error}") + if self.speaker_enabled: + t = Thread(target=self._receive_audio, daemon=True) + t.start() def on_error(self, error): self.logger.error(f"on_error: {error}") diff --git a/src/dailyai/services/local_stt_service.py b/src/dailyai/services/local_stt_service.py new file mode 100644 index 000000000..866d77cac --- /dev/null +++ b/src/dailyai/services/local_stt_service.py @@ -0,0 +1,72 @@ +import array +import io +import math +from typing import AsyncGenerator +import wave +from dailyai.queue_frame import AudioQueueFrame, QueueFrame, TextQueueFrame +from dailyai.services.ai_services import STTService + + +class LocalSTTService(STTService): + _content: io.BufferedRandom + _wave: wave.Wave_write + _current_silence_frames: int + + # Configuration + _min_rms: int + _max_silence_frames: int + _frame_rate: int + + def __init__(self, + min_rms: int = 400, + max_silence_frames: int = 3, + frame_rate: int = 16000, + **kwargs): + super().__init__(frame_rate, **kwargs) + self._current_silence_frames = 0 + self._min_rms = min_rms + self._max_silence_frames = max_silence_frames + self._frame_rate = frame_rate + self._new_wave() + + def _new_wave(self): + """Creates a new wave object and content buffer.""" + self._content = io.BufferedRandom(io.BytesIO()) + ww = wave.open(self._content, "wb") + ww.setnchannels(1) + ww.setsampwidth(2) + ww.setframerate(self._frame_rate) + self._wave = ww + + async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]: + """Processes a frame of audio data, either buffering or transcribing it.""" + if not isinstance(frame, AudioQueueFrame): + return + + data = frame.data + # Try to filter out empty background noise + # (Very rudimentary approach, can be improved) + rms = self._get_volume(data) + if rms >= self._min_rms: + # If volume is high enough, write new data to wave file + self._wave.writeframesraw(data) + + # If buffer is not empty and we detect a 3-frame pause in speech, + # transcribe the audio gathered so far. + if self._content.tell() > 0 and self._current_silence_frames > self._max_silence_frames: + self._current_silence_frames = 0 + self._wave.close() + self._content.seek(0) + text = await self.run_stt(self._content) + self._new_wave() + yield TextQueueFrame(text) + # If we get this far, this is a frame of silence + self._current_silence_frames += 1 + + def _get_volume(self, audio: bytes) -> float: + # https://docs.python.org/3/library/array.html + audio_array = array.array('h', audio) + squares = [sample**2 for sample in audio_array] + mean = sum(squares) / len(audio_array) + rms = math.sqrt(mean) + return rms diff --git a/src/dailyai/services/whisper_ai_services.py b/src/dailyai/services/whisper_ai_services.py new file mode 100644 index 000000000..88bb6f5d4 --- /dev/null +++ b/src/dailyai/services/whisper_ai_services.py @@ -0,0 +1,55 @@ +"""This module implements Whisper transcription with a locally-downloaded model.""" +import asyncio +from enum import Enum +import logging +from typing import BinaryIO +from faster_whisper import WhisperModel +from dailyai.services.local_stt_service import LocalSTTService + + +class Model(Enum): + """Class of basic Whisper model selection options""" + TINY = "tiny" + BASE = "base" + MEDIUM = "medium" + LARGE = "large-v3" + DISTIL_LARGE_V2 = "Systran/faster-distil-whisper-large-v2" + DISTIL_MEDIUM_EN = "Systran/faster-distil-whisper-medium.en" + + +class WhisperSTTService(LocalSTTService): + """Class to transcribe audio with a locally-downloaded Whisper model""" + _model: WhisperModel + + # Model configuration + _model_name: Model + _device: str + _compute_type: str + + def __init__(self, model_name: Model = Model.DISTIL_MEDIUM_EN, + device: str = "auto", + compute_type: str = "default"): + + super().__init__() + self.logger: logging.Logger = logging.getLogger("dailyai") + self._model_name = model_name + self._device = device + self._compute_type = compute_type + self._load() + + def _load(self): + """Loads the Whisper model. Note that if this is the first time + this model is being run, it will take time to download.""" + model = WhisperModel( + self._model_name.value, + device=self._device, + compute_type=self._compute_type) + self._model = model + + async def run_stt(self, audio: BinaryIO = None) -> str: + """Transcribes given audio using Whisper""" + segments, _ = await asyncio.to_thread(self._model.transcribe, audio) + res: str = "" + for segment in segments: + res += f"{segment.text} " + return res diff --git a/src/samples/foundational/07-whisper-transcription.py b/src/samples/foundational/07-whisper-transcription.py new file mode 100644 index 000000000..52376fc6c --- /dev/null +++ b/src/samples/foundational/07-whisper-transcription.py @@ -0,0 +1,45 @@ +import argparse +import asyncio +from threading import Thread + +from dailyai.services.daily_transport_service import DailyTransportService +from dailyai.services.whisper_ai_services import WhisperSTTService + + +async def main(room_url: str): + global transport + global stt + + transport = DailyTransportService( + room_url, + None, + "Transcription bot", + ) + transport.mic_enabled = False + transport.camera_enabled = False + transport.speaker_enabled = True + stt = WhisperSTTService() + transcription_output_queue = asyncio.Queue() + + async def handle_transcription(): + print("`````````TRANSCRIPTION`````````") + while True: + item = await transcription_output_queue.get() + print(item.text) + + async def handle_speaker(): + await stt.run_to_queue( + transcription_output_queue, + transport.get_receive_frames() + ) + await asyncio.gather(transport.run(), handle_speaker(), handle_transcription()) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Simple Daily Bot Sample") + parser.add_argument( + "-u", "--url", type=str, required=True, help="URL of the Daily room to join" + ) + + args, unknown = parser.parse_known_args() + asyncio.run(main(args.url))