From bfe2e0f36eb37d7a53808841a04ded40941ade77 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Wed, 10 Apr 2024 09:44:42 -0700
Subject: [PATCH] services: don't use image_size in ImageGenService

---
 src/dailyai/services/ai_services.py       |  5 ++---
 src/dailyai/services/azure_ai_services.py |  9 +++++----
 src/dailyai/services/fal_ai_services.py   |  4 ++--
 src/dailyai/services/open_ai_services.py  | 14 ++++++++------
 4 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py
index d4ecebee4..898c5f0a7 100644
--- a/src/dailyai/services/ai_services.py
+++ b/src/dailyai/services/ai_services.py
@@ -82,13 +82,12 @@ class TTSService(AIService):
 
 
 class ImageGenService(AIService):
-    def __init__(self, image_size, **kwargs):
+    def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.image_size = image_size
 
     # Renders the image. Returns an Image object.
     @abstractmethod
-    async def run_image_gen(self, sentence: str) -> tuple[str, bytes, tuple[int, int]]:
+    async def run_image_gen(self, prompt: str) -> tuple[str, bytes, tuple[int, int]]:
         pass
 
     async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
diff --git a/src/dailyai/services/azure_ai_services.py b/src/dailyai/services/azure_ai_services.py
index de068fbb4..151e78d62 100644
--- a/src/dailyai/services/azure_ai_services.py
+++ b/src/dailyai/services/azure_ai_services.py
@@ -97,23 +97,24 @@ class AzureImageGenServiceREST(ImageGenService):
         endpoint,
         model,
     ):
-        super().__init__(image_size=image_size)
+        super().__init__()
 
         self._api_key = api_key
         self._azure_endpoint = endpoint
         self._api_version = api_version
         self._model = model
         self._aiohttp_session = aiohttp_session
+        self._image_size = image_size
 
-    async def run_image_gen(self, sentence) -> tuple[str, bytes, tuple[int, int]]:
+    async def run_image_gen(self, prompt: str) -> 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,
             "Content-Type": "application/json"}
         body = {
             # Enter your prompt text here
-            "prompt": sentence,
-            "size": self.image_size,
+            "prompt": prompt,
+            "size": self._image_size,
             "n": 1,
         }
         async with self._aiohttp_session.post(
diff --git a/src/dailyai/services/fal_ai_services.py b/src/dailyai/services/fal_ai_services.py
index 8e6833ee8..5c1d15151 100644
--- a/src/dailyai/services/fal_ai_services.py
+++ b/src/dailyai/services/fal_ai_services.py
@@ -45,7 +45,7 @@ class FalImageGenService(ImageGenService):
         if key_secret:
             os.environ["FAL_KEY_SECRET"] = key_secret
 
-    async def run_image_gen(self, prompt) -> tuple[str, bytes]:
+    async def run_image_gen(self, prompt: str) -> tuple[str, bytes, tuple[int, int]]:
         def get_image_url(prompt):
             handler = fal.apps.submit(  # type: ignore
                 self._model,
@@ -72,4 +72,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)
diff --git a/src/dailyai/services/open_ai_services.py b/src/dailyai/services/open_ai_services.py
index 95045f6dc..9eaec5218 100644
--- a/src/dailyai/services/open_ai_services.py
+++ b/src/dailyai/services/open_ai_services.py
@@ -1,3 +1,4 @@
+from typing import Literal
 import aiohttp
 from PIL import Image
 import io
@@ -26,24 +27,25 @@ class OpenAIImageGenService(ImageGenService):
     def __init__(
         self,
         *,
-        image_size: str,
+        image_size: Literal["256x256", "512x512", "1024x1024", "1792x1024", "1024x1792"],
         aiohttp_session: aiohttp.ClientSession,
         api_key,
         model="dall-e-3",
     ):
-        super().__init__(image_size=image_size)
+        super().__init__()
         self._model = model
+        self._image_size = image_size
         self._client = AsyncOpenAI(api_key=api_key)
         self._aiohttp_session = aiohttp_session
 
-    async def run_image_gen(self, sentence) -> tuple[str, bytes, tuple[int, int]]:
-        self.logger.info("Generating OpenAI image", sentence)
+    async def run_image_gen(self, prompt: str) -> tuple[str, bytes, tuple[int, int]]:
+        self.logger.info("Generating OpenAI image", prompt)
 
         image = await self._client.images.generate(
-            prompt=sentence,
+            prompt=prompt,
             model=self._model,
             n=1,
-            size=self.image_size
+            size=self._image_size
         )
         image_url = image.data[0].url
         if not image_url: