59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
import aiohttp
|
|
import asyncio
|
|
import io
|
|
import os
|
|
from PIL import Image
|
|
|
|
from dailyai.services.ai_services import ImageGenService
|
|
|
|
try:
|
|
import fal
|
|
except ModuleNotFoundError as e:
|
|
print(f"Exception: {e}")
|
|
print(
|
|
"In order to use Fal, you need to `pip install dailyai[fal]`. Also, set `FAL_KEY_ID` and `FAL_KEY_SECRET` environment variables.")
|
|
raise Exception(f"Missing module: {e}")
|
|
|
|
|
|
class FalImageGenService(ImageGenService):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
image_size,
|
|
aiohttp_session: aiohttp.ClientSession,
|
|
key_id=None,
|
|
key_secret=None
|
|
):
|
|
super().__init__(image_size)
|
|
self._aiohttp_session = aiohttp_session
|
|
if key_id:
|
|
os.environ["FAL_KEY_ID"] = key_id
|
|
if key_secret:
|
|
os.environ["FAL_KEY_SECRET"] = key_secret
|
|
|
|
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",
|
|
# "fal-ai/fast-sdxl",
|
|
arguments={"prompt": sentence},
|
|
)
|
|
for event in handler.iter_events():
|
|
if isinstance(event, fal.apps.InProgress):
|
|
pass
|
|
|
|
result = handler.get()
|
|
|
|
image_url = result["images"][0]["url"] if result else None
|
|
if not image_url:
|
|
raise Exception("Image generation failed")
|
|
|
|
return image_url
|
|
|
|
image_url = await asyncio.to_thread(get_image_url, sentence, self.image_size)
|
|
# Load the image from the url
|
|
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(), image.size)
|