wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

updated ImageFrame and added URLImageFrame and UserImageFrame

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

7
src/dailyai/services/ai_services.py
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):

4
src/dailyai/services/azure_ai_services.py
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)

4
src/dailyai/services/fal_ai_services.py
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)

4
src/dailyai/services/open_ai_services.py
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)

2
src/dailyai/services/to_be_updated/mock_ai_service.py
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):