504 lines
18 KiB
Python
504 lines
18 KiB
Python
from abc import abstractmethod
|
|
import asyncio
|
|
import itertools
|
|
import numpy as np
|
|
|
|
import queue
|
|
import threading
|
|
import time
|
|
from typing import Any, AsyncGenerator
|
|
from enum import Enum
|
|
from dailyai.pipeline.frame_processor import FrameProcessor
|
|
|
|
from dailyai.pipeline.frames import (
|
|
SendAppMessageFrame,
|
|
AudioFrame,
|
|
EndFrame,
|
|
ImageFrame,
|
|
Frame,
|
|
PipelineStartedFrame,
|
|
SpriteFrame,
|
|
StartFrame,
|
|
TextFrame,
|
|
UserImageRequestFrame,
|
|
UserStartedSpeakingFrame,
|
|
UserStoppedSpeakingFrame,
|
|
)
|
|
from dailyai.pipeline.pipeline import Pipeline
|
|
from dailyai.services.ai_services import TTSService
|
|
from dailyai.transports.abstract_transport import AbstractTransport
|
|
|
|
|
|
# Provided by Alexander Veysov
|
|
|
|
|
|
def int2float(sound):
|
|
try:
|
|
abs_max = np.abs(sound).max()
|
|
sound = sound.astype("float32")
|
|
if abs_max > 0:
|
|
sound *= 1 / 32768
|
|
sound = sound.squeeze() # depends on the use case
|
|
return sound
|
|
except ValueError:
|
|
return sound
|
|
|
|
|
|
class VADState(Enum):
|
|
QUIET = 1
|
|
STARTING = 2
|
|
SPEAKING = 3
|
|
STOPPING = 4
|
|
|
|
|
|
class ThreadedTransport(AbstractTransport):
|
|
|
|
def __init__(
|
|
self,
|
|
**kwargs,
|
|
) -> None:
|
|
super().__init__(**kwargs)
|
|
self._vad_start_s = kwargs.get("vad_start_s") or 0.2
|
|
self._vad_stop_s = kwargs.get("vad_stop_s") or 0.8
|
|
self._context = kwargs.get("context") or []
|
|
self._vad_enabled = kwargs.get("vad_enabled") or False
|
|
self._has_webrtc_vad = kwargs.get("has_webrtc_vad") or False
|
|
if self._vad_enabled and self._speaker_enabled:
|
|
raise Exception(
|
|
"Sorry, you can't use speaker_enabled and vad_enabled at the same time. Please set one to False."
|
|
)
|
|
self._vad_samples = 1536
|
|
|
|
if self._vad_enabled:
|
|
try:
|
|
global torch, torchaudio
|
|
import torch
|
|
# We don't use torchaudio here, but we need to try importing it because
|
|
# Silero uses it
|
|
import torchaudio
|
|
torch.set_num_threads(1)
|
|
|
|
(self.model, self.utils) = torch.hub.load(
|
|
repo_or_dir="snakers4/silero-vad", model="silero_vad", force_reload=False
|
|
)
|
|
self._logger.debug("Loaded Silero VAD")
|
|
|
|
except ModuleNotFoundError as e:
|
|
if self._has_webrtc_vad:
|
|
self._logger.debug(
|
|
f"Couldn't load torch; using webrtc VAD")
|
|
self._vad_samples = int(self._speaker_sample_rate / 100.0)
|
|
else:
|
|
self._logger.error(f"Exception: {e}")
|
|
self._logger.error(
|
|
"In order to use Silero VAD, you'll need to `pip install dailyai[silero].")
|
|
raise Exception(f"Missing module(s): {e}")
|
|
|
|
vad_frame_s = self._vad_samples / self._speaker_sample_rate
|
|
self._vad_start_frames = round(self._vad_start_s / vad_frame_s)
|
|
self._vad_stop_frames = round(self._vad_stop_s / vad_frame_s)
|
|
self._vad_starting_count = 0
|
|
self._vad_stopping_count = 0
|
|
self._vad_state = VADState.QUIET
|
|
self._user_is_speaking = False
|
|
|
|
self._threadsafe_send_queue = queue.Queue()
|
|
|
|
self._images = None
|
|
|
|
try:
|
|
self._loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
self._loop = None
|
|
|
|
self._stop_threads = threading.Event()
|
|
self._is_interrupted = threading.Event()
|
|
|
|
async def run(self, pipeline: Pipeline | None = None, override_pipeline_source_queue=True):
|
|
self._prerun()
|
|
|
|
async_output_queue_marshal_task = asyncio.create_task(
|
|
self._marshal_frames())
|
|
|
|
self._camera_thread = threading.Thread(
|
|
target=self._run_camera, daemon=True)
|
|
self._camera_thread.start()
|
|
|
|
self._frame_consumer_thread = threading.Thread(
|
|
target=self._frame_consumer, daemon=True
|
|
)
|
|
self._frame_consumer_thread.start()
|
|
|
|
if self._speaker_enabled:
|
|
self._receive_audio_thread = threading.Thread(
|
|
target=self._receive_audio, daemon=True
|
|
)
|
|
self._receive_audio_thread.start()
|
|
|
|
if self._vad_enabled:
|
|
self._vad_thread = threading.Thread(target=self._vad, daemon=True)
|
|
self._vad_thread.start()
|
|
|
|
pipeline_task = None
|
|
if pipeline:
|
|
pipeline_task = asyncio.create_task(
|
|
self.run_pipeline(pipeline, override_pipeline_source_queue)
|
|
)
|
|
|
|
try:
|
|
while time.time() < self._expiration and not self._stop_threads.is_set():
|
|
await asyncio.sleep(1)
|
|
except Exception as e:
|
|
self._logger.error(f"Exception {e}")
|
|
raise e
|
|
finally:
|
|
# Do anything that must be done to clean up
|
|
self._post_run()
|
|
|
|
self._stop_threads.set()
|
|
|
|
if pipeline_task:
|
|
pipeline_task.cancel()
|
|
|
|
await self.send_queue.put(EndFrame())
|
|
|
|
await async_output_queue_marshal_task
|
|
self._frame_consumer_thread.join()
|
|
|
|
if self._speaker_enabled:
|
|
self._receive_audio_thread.join()
|
|
|
|
if self._vad_enabled:
|
|
self._vad_thread.join()
|
|
|
|
async def run_pipeline(self, pipeline: Pipeline, override_pipeline_source_queue=True):
|
|
pipeline.set_sink(self.send_queue)
|
|
if override_pipeline_source_queue:
|
|
pipeline.set_source(self.receive_queue)
|
|
await pipeline.run_pipeline()
|
|
|
|
async def run_interruptible_pipeline(
|
|
self,
|
|
pipeline: Pipeline,
|
|
pre_processor: FrameProcessor | None = None,
|
|
post_processor: FrameProcessor | None = None,
|
|
):
|
|
pipeline.set_sink(self.send_queue)
|
|
source_queue = asyncio.Queue()
|
|
pipeline.set_source(source_queue)
|
|
pipeline_task = asyncio.create_task(pipeline.run_pipeline())
|
|
|
|
async def yield_frame(frame: Frame) -> AsyncGenerator[Frame, None]:
|
|
yield frame
|
|
|
|
async def post_process(post_processor: FrameProcessor):
|
|
while True:
|
|
frame = await self.completed_queue.get()
|
|
|
|
# We ignore the output of the post_processor's process frame;
|
|
# this is called to update the post-processor's state.
|
|
async for frame in post_processor.process_frame(frame):
|
|
pass
|
|
|
|
if isinstance(frame, EndFrame):
|
|
break
|
|
|
|
if post_processor:
|
|
post_process_task = asyncio.create_task(
|
|
post_process(post_processor))
|
|
|
|
started = False
|
|
|
|
async for frame in self.get_receive_frames():
|
|
if isinstance(frame, UserStartedSpeakingFrame):
|
|
pipeline_task.cancel()
|
|
self.interrupt()
|
|
pipeline_task = asyncio.create_task(pipeline.run_pipeline())
|
|
started = False
|
|
|
|
if not started:
|
|
await self.send_queue.put(StartFrame())
|
|
|
|
if pre_processor:
|
|
frame_generator = pre_processor.process_frame(frame)
|
|
else:
|
|
frame_generator = yield_frame(frame)
|
|
|
|
async for frame in frame_generator:
|
|
await source_queue.put(frame)
|
|
|
|
if isinstance(frame, EndFrame):
|
|
break
|
|
|
|
await asyncio.gather(pipeline_task, post_process_task)
|
|
|
|
async def say(self, text: str, tts: TTSService):
|
|
"""Say a phrase. Use with caution; this bypasses any running pipelines."""
|
|
async for frame in tts.process_frame(TextFrame(text)):
|
|
await self.send_queue.put(frame)
|
|
|
|
def _post_run(self):
|
|
# Note that this function must be idempotent! It can be called multiple times
|
|
# if, for example, a keyboard interrupt occurs.
|
|
pass
|
|
|
|
def stop(self):
|
|
self._stop_threads.set()
|
|
|
|
async def stop_when_done(self):
|
|
await self._wait_for_send_queue_to_empty()
|
|
self.stop()
|
|
|
|
async def _wait_for_send_queue_to_empty(self):
|
|
await self.send_queue.join()
|
|
self._threadsafe_send_queue.join()
|
|
|
|
@abstractmethod
|
|
def write_frame_to_camera(self, frame: bytes):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def write_frame_to_mic(self, frame: bytes):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def read_audio_frames(self, desired_frame_count):
|
|
return bytes()
|
|
|
|
@abstractmethod
|
|
def _prerun(self):
|
|
pass
|
|
|
|
def _silero_vad_analyze(self):
|
|
try:
|
|
audio_chunk = self.read_audio_frames(self._vad_samples)
|
|
audio_int16 = np.frombuffer(audio_chunk, np.int16)
|
|
audio_float32 = int2float(audio_int16)
|
|
new_confidence = self.model(
|
|
torch.from_numpy(audio_float32), 16000).item()
|
|
# yeses = int(new_confidence * 20.0)
|
|
# nos = 20 - yeses
|
|
# out = "!" * yeses + "." * nos
|
|
# print(f"!!! confidence: {out}")
|
|
speaking = new_confidence > 0.5
|
|
return speaking
|
|
except BaseException:
|
|
# This comes from an empty audio array
|
|
return False
|
|
|
|
def _vad(self):
|
|
|
|
while not self._stop_threads.is_set():
|
|
if hasattr(self, 'model'): # we can use Silero
|
|
speaking = self._silero_vad_analyze()
|
|
elif self._has_webrtc_vad:
|
|
speaking = self._webrtc_vad_analyze()
|
|
else:
|
|
raise Exception("VAD is running with no VAD service available")
|
|
if speaking:
|
|
match self._vad_state:
|
|
case VADState.QUIET:
|
|
self._vad_state = VADState.STARTING
|
|
self._vad_starting_count = 1
|
|
case VADState.STARTING:
|
|
self._vad_starting_count += 1
|
|
case VADState.STOPPING:
|
|
self._vad_state = VADState.SPEAKING
|
|
self._vad_stopping_count = 0
|
|
else:
|
|
match self._vad_state:
|
|
case VADState.STARTING:
|
|
self._vad_state = VADState.QUIET
|
|
self._vad_starting_count = 0
|
|
case VADState.SPEAKING:
|
|
self._vad_state = VADState.STOPPING
|
|
self._vad_stopping_count = 1
|
|
case VADState.STOPPING:
|
|
self._vad_stopping_count += 1
|
|
|
|
if (
|
|
self._vad_state == VADState.STARTING
|
|
and self._vad_starting_count >= self._vad_start_frames
|
|
):
|
|
if self._loop:
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.receive_queue.put(
|
|
UserStartedSpeakingFrame()), self._loop)
|
|
# self.interrupt()
|
|
self._vad_state = VADState.SPEAKING
|
|
self._vad_starting_count = 0
|
|
if (
|
|
self._vad_state == VADState.STOPPING
|
|
and self._vad_stopping_count >= self._vad_stop_frames
|
|
):
|
|
|
|
if self._loop:
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.receive_queue.put(
|
|
UserStoppedSpeakingFrame()), self._loop)
|
|
self._vad_state = VADState.QUIET
|
|
self._vad_stopping_count = 0
|
|
|
|
async def _marshal_frames(self):
|
|
while True:
|
|
frame: Frame | list = await self.send_queue.get()
|
|
self._threadsafe_send_queue.put(frame)
|
|
self.send_queue.task_done()
|
|
if isinstance(frame, EndFrame):
|
|
break
|
|
|
|
def interrupt(self):
|
|
self._logger.debug("### Interrupting")
|
|
self._is_interrupted.set()
|
|
|
|
async def get_receive_frames(self) -> AsyncGenerator[Frame, None]:
|
|
while True:
|
|
frame = await self.receive_queue.get()
|
|
yield frame
|
|
if isinstance(frame, EndFrame):
|
|
break
|
|
|
|
def _receive_audio(self):
|
|
if not self._loop:
|
|
self._logger.error("No loop available for audio thread")
|
|
return
|
|
|
|
seconds = 1
|
|
desired_frame_count = self._speaker_sample_rate * seconds
|
|
while not self._stop_threads.is_set():
|
|
buffer = self.read_audio_frames(desired_frame_count)
|
|
if len(buffer) > 0:
|
|
frame = AudioFrame(buffer)
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.receive_queue.put(frame), self._loop
|
|
)
|
|
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.receive_queue.put(
|
|
EndFrame()), self._loop)
|
|
|
|
def _set_image(self, image: bytes):
|
|
self._images = itertools.cycle([image])
|
|
|
|
def _set_images(self, images: list[bytes], start_frame=0):
|
|
self._images = itertools.cycle(images)
|
|
|
|
def request_participant_image(self, participant_id: str):
|
|
""" Child classes should override this to force an image from a user. """
|
|
pass
|
|
|
|
def send_app_message(self, message: Any, participant_id: str | None):
|
|
""" Child classes should override this to send a custom message to the room. """
|
|
pass
|
|
|
|
def _run_camera(self):
|
|
try:
|
|
while not self._stop_threads.is_set():
|
|
if self._images:
|
|
this_frame = next(self._images)
|
|
self.write_frame_to_camera(this_frame)
|
|
|
|
time.sleep(1.0 / self._camera_framerate)
|
|
except Exception as e:
|
|
self._logger.error(f"Exception {e} in camera thread.")
|
|
raise e
|
|
|
|
def _frame_consumer(self):
|
|
self._logger.info("🎬 Starting frame consumer thread")
|
|
b = bytearray()
|
|
smallest_write_size = 3200
|
|
largest_write_size = 8000
|
|
while True:
|
|
try:
|
|
frames_or_frame: Frame | list[Frame] = self._threadsafe_send_queue.get(
|
|
)
|
|
if (
|
|
isinstance(frames_or_frame, AudioFrame)
|
|
and len(frames_or_frame.data) > largest_write_size
|
|
):
|
|
# subdivide large audio frames to enable interruption
|
|
frames = []
|
|
for i in range(0, len(frames_or_frame.data),
|
|
largest_write_size):
|
|
frames.append(AudioFrame(
|
|
frames_or_frame.data[i: i + largest_write_size]))
|
|
elif isinstance(frames_or_frame, Frame):
|
|
frames: list[Frame] = [frames_or_frame]
|
|
elif isinstance(frames_or_frame, list):
|
|
frames: list[Frame] = frames_or_frame
|
|
else:
|
|
raise Exception("Unknown type in output queue")
|
|
|
|
for frame in frames:
|
|
if isinstance(frame, EndFrame):
|
|
self._logger.info("Stopping frame consumer thread")
|
|
self._stop_threads.set()
|
|
self._threadsafe_send_queue.task_done()
|
|
if self._loop:
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.completed_queue.put(frame), self._loop
|
|
)
|
|
# Also send the EndFrame to the pipeline so it can stop
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.receive_queue.put(frame), self._loop
|
|
)
|
|
return
|
|
|
|
# if interrupted, we just pull frames off the queue and
|
|
# discard them
|
|
if not self._is_interrupted.is_set():
|
|
if frame:
|
|
if isinstance(frame, AudioFrame):
|
|
chunk = frame.data
|
|
|
|
b.extend(chunk)
|
|
truncated_length: int = len(b) - (
|
|
len(b) % smallest_write_size
|
|
)
|
|
if truncated_length:
|
|
self.write_frame_to_mic(
|
|
bytes(b[:truncated_length]))
|
|
b = b[truncated_length:]
|
|
elif isinstance(frame, ImageFrame):
|
|
self._set_image(frame.image)
|
|
elif isinstance(frame, SpriteFrame):
|
|
self._set_images(frame.images)
|
|
elif isinstance(frame, UserImageRequestFrame):
|
|
self.request_participant_image(frame.user_id)
|
|
elif isinstance(frame, SendAppMessageFrame):
|
|
self.send_app_message(frame.message, frame.participant_id)
|
|
elif len(b):
|
|
self.write_frame_to_mic(bytes(b))
|
|
b = bytearray()
|
|
else:
|
|
# if there are leftover audio bytes, write them now; failing to do so
|
|
# can cause static in the audio stream.
|
|
if len(b):
|
|
truncated_length = len(b) - (len(b) % 160)
|
|
self.write_frame_to_mic(
|
|
bytes(b[:truncated_length]))
|
|
b = bytearray()
|
|
|
|
if isinstance(frame, StartFrame):
|
|
self._is_interrupted.clear()
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.receive_queue.put(PipelineStartedFrame()),
|
|
self._loop,
|
|
)
|
|
|
|
if self._loop:
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.completed_queue.put(frame), self._loop
|
|
)
|
|
|
|
self._threadsafe_send_queue.task_done()
|
|
except queue.Empty:
|
|
if len(b):
|
|
self.write_frame_to_mic(bytes(b))
|
|
|
|
b = bytearray()
|
|
except Exception as e:
|
|
self._logger.error(
|
|
f"Exception in frame_consumer: {e}, {len(b)}")
|
|
raise e
|