changed fal to a maybe slightly faster model

This commit is contained in:
Kwindla Hultman Kramer
2024-03-10 19:40:51 -07:00
parent ef39d842a5
commit 37e904ce68

View File

@@ -9,17 +9,19 @@ from dailyai.services.ai_services import ImageGenService
from dailyai.services.ai_services import ImageGenService from dailyai.services.ai_services import ImageGenService
# Fal expects FAL_KEY_ID and FAL_KEY_SECRET to be set in the env # Fal expects FAL_KEY_ID and FAL_KEY_SECRET to be set in the env
class FalImageGenService(ImageGenService): class FalImageGenService(ImageGenService):
def __init__( def __init__(
self, self,
*, *,
image_size, image_size,
aiohttp_session: aiohttp.ClientSession, aiohttp_session: aiohttp.ClientSession,
key_id=None, key_id=None,
key_secret=None): key_secret=None
):
super().__init__(image_size) super().__init__(image_size)
self._aiohttp_session = aiohttp_session self._aiohttp_session = aiohttp_session
if key_id: if key_id:
@@ -30,10 +32,9 @@ class FalImageGenService(ImageGenService):
async def run_image_gen(self, sentence) -> tuple[str, bytes]: async def run_image_gen(self, sentence) -> tuple[str, bytes]:
def get_image_url(sentence, size): def get_image_url(sentence, size):
handler = fal.apps.submit( handler = fal.apps.submit(
"110602490-fast-sdxl", # "110602490-fast-sdxl",
arguments={ "fal-ai/fast-sdxl",
"prompt": sentence arguments={"prompt": sentence},
},
) )
for event in handler.iter_events(): for event in handler.iter_events():
if isinstance(event, fal.apps.InProgress): if isinstance(event, fal.apps.InProgress):
@@ -46,6 +47,7 @@ class FalImageGenService(ImageGenService):
raise Exception("Image generation failed") raise Exception("Image generation failed")
return image_url return image_url
image_url = await asyncio.to_thread(get_image_url, sentence, self.image_size) image_url = await asyncio.to_thread(get_image_url, sentence, self.image_size)
# Load the image from the url # Load the image from the url
async with self._aiohttp_session.get(image_url) as response: async with self._aiohttp_session.get(image_url) as response: