This commit is contained in:
Moishe Lettvin
2024-01-25 14:55:51 -05:00
parent 6da78dbf9c
commit 42c142aff0
2 changed files with 32 additions and 33 deletions

View File

@@ -101,39 +101,37 @@ class AzureImageGenServiceREST(ImageGenService):
self.aiohttp_client_session = aiohttp_client_session or aiohttp.ClientSession() self.aiohttp_client_session = aiohttp_client_session or aiohttp.ClientSession()
async def run_image_gen(self, sentence) -> tuple[str, bytes]: async def run_image_gen(self, sentence) -> tuple[str, bytes]:
# TODO hoist the session to app-level url = f"{self.azure_endpoint}openai/images/generations:submit?api-version={self.api_version}"
async with self.aiohttp_client_session as session: headers= { "api-key": self.api_key, "Content-Type": "application/json" }
url = f"{self.azure_endpoint}openai/images/generations:submit?api-version={self.api_version}" body = {
headers= { "api-key": self.api_key, "Content-Type": "application/json" } "prompt": sentence,
body = { "size": self.image_size,
"prompt": sentence, "n": 1,
"size": self.image_size, }
"n": 1, async with self.aiohttp_client_session.post(
} url, headers=headers, json=body
async with self.aiohttp_client_session.post( ) as submission:
url, headers=headers, json=body operation_location = submission.headers['operation-location']
) as submission:
operation_location = submission.headers['operation-location']
status = "" status = ""
attempts_left = 120 attempts_left = 120
json_response = None json_response = None
while status != "succeeded": while status != "succeeded":
attempts_left -= 1 attempts_left -= 1
if attempts_left == 0: if attempts_left == 0:
raise Exception("Image generation timed out") raise Exception("Image generation timed out")
await asyncio.sleep(1) await asyncio.sleep(1)
response = await session.get(operation_location, headers=headers) response = await self.aiohttp_client_session.get(operation_location, headers=headers)
json_response = await response.json() json_response = await response.json()
status = json_response["status"] status = json_response["status"]
image_url = json_response["result"]["data"][0]["url"] if json_response else None image_url = json_response["result"]["data"][0]["url"] if json_response else None
if not image_url: if not image_url:
raise Exception("Image generation failed") raise Exception("Image generation failed")
# Load the image from the url # Load the image from the url
async with session.get(image_url) as response: async with self.aiohttp_client_session.get(image_url) as response:
image_stream = io.BytesIO(await response.content.read()) image_stream = io.BytesIO(await response.content.read())
image = Image.open(image_stream) image = Image.open(image_stream)
return (image_url, image.tobytes()) return (image_url, image.tobytes())

View File

@@ -4,6 +4,7 @@ import asyncio
from dailyai.queue_frame import TextQueueFrame from dailyai.queue_frame import TextQueueFrame
from dailyai.services.daily_transport_service import DailyTransportService from dailyai.services.daily_transport_service import DailyTransportService
from dailyai.services.open_ai_services import OpenAIImageGenService from dailyai.services.open_ai_services import OpenAIImageGenService
from dailyai.services.azure_ai_services import AzureImageGenServiceREST
local_joined = False local_joined = False
participant_joined = False participant_joined = False
@@ -21,7 +22,7 @@ async def main(room_url):
transport.camera_width = 1024 transport.camera_width = 1024
transport.camera_height = 1024 transport.camera_height = 1024
imagegen = OpenAIImageGenService(image_size="1024x1024") imagegen = AzureImageGenServiceREST(image_size="1024x1024")
image_task = asyncio.create_task( image_task = asyncio.create_task(
imagegen.run_to_queue(transport.send_queue, [TextQueueFrame("a cat in the style of picasso")]) imagegen.run_to_queue(transport.send_queue, [TextQueueFrame("a cat in the style of picasso")])
) )