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

View File

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