updated ImageFrame and added URLImageFrame and UserImageFrame

This commit is contained in:
Aleix Conchillo Flaqué
2024-04-08 23:10:10 -07:00
parent 96950ca6df
commit 319b8e7816
16 changed files with 97 additions and 30 deletions

View File

@@ -360,7 +360,7 @@ class GatedAggregator(FrameProcessor):
... start_open=False)
>>> asyncio.run(print_frames(aggregator, TextFrame("Hello")))
>>> asyncio.run(print_frames(aggregator, TextFrame("Hello again.")))
>>> asyncio.run(print_frames(aggregator, ImageFrame(url='', image=bytes([]))))
>>> asyncio.run(print_frames(aggregator, ImageFrame(image=bytes([]), size=(0, 0))))
ImageFrame
Hello
Hello again.

View File

@@ -70,11 +70,39 @@ class AudioFrame(Frame):
class ImageFrame(Frame):
"""An image. Will be shown by the transport if the transport's camera is
enabled."""
url: str | None
image: bytes
size: tuple[int, int]
def __str__(self):
return f"{self.__class__.__name__}, url: {self.url}, image size: {len(self.image)} B"
return f"{self.__class__.__name__}, image size: {self.size[0]}x{self.size[1]} buffer size: {len(self.image)} B"
@dataclass()
class URLImageFrame(ImageFrame):
"""An image. Will be shown by the transport if the transport's camera is
enabled."""
url: str | None
def __init__(self, url, image, size):
super().__init__(image, size)
self.url = url
def __str__(self):
return f"{self.__class__.__name__}, url: {self.url}, image size: {self.size[0]}x{self.size[1]}, buffer size: {len(self.image)} B"
@dataclass()
class UserImageFrame(ImageFrame):
"""An image associated to a user. Will be shown by the transport if the transport's camera is
enabled."""
user_id: str
def __init__(self, user_id, image, size):
super().__init__(image, size)
self.user_id = user_id
def __str__(self):
return f"{self.__class__.__name__}, user: {self.user_id}, image size: {self.size[0]}x{self.size[1]}, buffer size: {len(self.image)} B"
@dataclass()

View File

@@ -14,6 +14,7 @@ from dailyai.pipeline.frames import (
TTSStartFrame,
TextFrame,
TranscriptionFrame,
URLImageFrame,
)
from abc import abstractmethod
@@ -87,7 +88,7 @@ class ImageGenService(AIService):
# Renders the image. Returns an Image object.
@abstractmethod
async def run_image_gen(self, sentence: str) -> tuple[str, bytes]:
async def run_image_gen(self, sentence: str) -> tuple[str, bytes, tuple[int, int]]:
pass
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
@@ -95,8 +96,8 @@ class ImageGenService(AIService):
yield frame
return
(url, image_data) = await self.run_image_gen(frame.text)
yield ImageFrame(url, image_data)
(url, image_data, image_size) = await self.run_image_gen(frame.text)
yield URLImageFrame(url, image_data, image_size)
class STTService(AIService):

View File

@@ -105,7 +105,7 @@ class AzureImageGenServiceREST(ImageGenService):
self._model = model
self._aiohttp_session = aiohttp_session
async def run_image_gen(self, sentence) -> tuple[str, bytes]:
async def run_image_gen(self, sentence) -> tuple[str, bytes, tuple[int, int]]:
url = f"{self._azure_endpoint}openai/images/generations:submit?api-version={self._api_version}"
headers = {
"api-key": self._api_key,
@@ -146,4 +146,4 @@ class AzureImageGenServiceREST(ImageGenService):
async with self._aiohttp_session.get(image_url) as response:
image_stream = io.BytesIO(await response.content.read())
image = Image.open(image_stream)
return (image_url, image.tobytes())
return (image_url, image.tobytes(), image.size)

View File

@@ -31,7 +31,7 @@ class FalImageGenService(ImageGenService):
if key_secret:
os.environ["FAL_KEY_SECRET"] = key_secret
async def run_image_gen(self, sentence) -> tuple[str, bytes]:
async def run_image_gen(self, sentence) -> tuple[str, bytes, tuple[int, int]]:
def get_image_url(sentence, size):
handler = fal.apps.submit(
"110602490-fast-sdxl",
@@ -55,4 +55,4 @@ class FalImageGenService(ImageGenService):
async with self._aiohttp_session.get(image_url) as response:
image_stream = io.BytesIO(await response.content.read())
image = Image.open(image_stream)
return (image_url, image.tobytes())
return (image_url, image.tobytes(), image.size)

View File

@@ -36,7 +36,7 @@ class OpenAIImageGenService(ImageGenService):
self._client = AsyncOpenAI(api_key=api_key)
self._aiohttp_session = aiohttp_session
async def run_image_gen(self, sentence) -> tuple[str, bytes]:
async def run_image_gen(self, sentence) -> tuple[str, bytes, tuple[int, int]]:
self.logger.info("Generating OpenAI image", sentence)
image = await self._client.images.generate(
@@ -53,4 +53,4 @@ class OpenAIImageGenService(ImageGenService):
async with self._aiohttp_session.get(image_url) as response:
image_stream = io.BytesIO(await response.content.read())
image = Image.open(image_stream)
return (image_url, image.tobytes())
return (image_url, image.tobytes(), image.size)

View File

@@ -19,7 +19,7 @@ class MockAIService(AIService):
image_stream = io.BytesIO(response.content)
image = Image.open(image_stream)
time.sleep(1)
return (image_url, image)
return (image_url, image.tobytes(), image.size)
def run_llm(self, messages, latest_user_message=None, stream=True):
for i in range(5):

View File

@@ -2,6 +2,7 @@ import asyncio
import inspect
import logging
import signal
import time
import threading
import types
@@ -11,6 +12,7 @@ from typing import Any
from dailyai.pipeline.frames import (
ReceivedAppMessageFrame,
TranscriptionFrame,
UserImageFrame,
)
from threading import Event
@@ -58,6 +60,7 @@ class DailyTransport(ThreadedTransport, EventHandler):
bot_name: str,
min_others_count: int = 1,
start_transcription: bool = False,
video_rendering_enabled: bool = False,
**kwargs,
):
kwargs['has_webrtc_vad'] = True
@@ -69,6 +72,7 @@ class DailyTransport(ThreadedTransport, EventHandler):
self._token: str | None = token
self._min_others_count = min_others_count
self._start_transcription = start_transcription
self._video_rendering_enabled = video_rendering_enabled
self._is_interrupted = Event()
self._stop_threads = Event()
@@ -76,6 +80,8 @@ class DailyTransport(ThreadedTransport, EventHandler):
self._other_participant_has_joined = False
self._my_participant_id = None
self._video_renderers = {}
self.transcription_settings = {
"language": "en",
"tier": "nova",
@@ -236,7 +242,7 @@ class DailyTransport(ThreadedTransport, EventHandler):
self.client.update_subscription_profiles({
"base": {
"camera": "unsubscribed",
"camera": "subscribed" if self._video_rendering_enabled else "unsubscribed",
}
})
@@ -268,6 +274,37 @@ class DailyTransport(ThreadedTransport, EventHandler):
def start_recording(self):
self.client.start_recording()
def render_participant_video(self,
participant_id,
framerate=10,
video_source="camera",
color_format="RGB") -> None:
if not self._video_rendering_enabled:
self._logger.warn("Video rendering is not enabled")
self._video_renderers[participant_id] = {
"framerate": framerate,
"timestamp": 0,
}
self.client.set_video_renderer(
participant_id,
self.on_participant_video_frame,
video_source=video_source,
color_format=color_format)
def on_participant_video_frame(self, participant_id, video_frame):
curr_time = time.time()
prev_time = self._video_renderers[participant_id]["timestamp"]
diff_time = curr_time - prev_time
period = 1 / self._video_renderers[participant_id]["framerate"]
if diff_time > period and self._loop:
self._video_renderers[participant_id]["timestamp"] = curr_time
frame = UserImageFrame(participant_id, video_frame.buffer,
(video_frame.width, video_frame.height))
asyncio.run_coroutine_threadsafe(
self.receive_queue.put(frame), self._loop
)
def on_error(self, error):
self._logger.error(f"on_error: {error}")