Changes.

2024-03-16 13:44:46 +01:00
parent 39d7c48ee3
commit 0c23da7c9e
7 changed files with 297 additions and 163 deletions
@@ -0,0 +1,66 @@
+import logging
+import asyncio
+import aiohttp
+from .ai_responder import exponential_backoff, AIResponderBase
+from io import BytesIO
+
+
+class LeonardoAIDrawMixIn(AIResponderBase):
+    async def draw_leonardo(self, description: str) -> BytesIO:
+        error_backoff = exponential_backoff(max_attempts=12)
+        generation_id = None
+        image_url = None
+        image_bytes = None
+        while True:
+            error_sleep = next(error_backoff)
+            try:
+                async with aiohttp.ClientSession() as session:
+                    if generation_id is None:
+                        async with session.post("https://cloud.leonardo.ai/api/rest/v1/generations",
+                                                json={"prompt": description,
+                                                      "modelId": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3",
+                                                      "num_images": 1,
+                                                      "sd_version": "v2",
+                                                      "promptMagic": True,
+                                                      "unzoomAmount": 1,
+                                                      "width": 512,
+                                                      "height": 512},
+                                                headers={"Authorization": f"Bearer {self.config['leonardo-token']}",
+                                                         "Accept": "application/json",
+                                                         "Content-Type": "application/json"},
+                                                ) as response:
+                            response = await response.json()
+                            if "sdGenerationJob" not in response:
+                                logging.warning(f"No 'sdGenerationJob' found in response, sleep for {error_sleep}s: {repr(response)}")
+                                await asyncio.sleep(error_sleep)
+                                continue
+                            generation_id = response["sdGenerationJob"]["generationId"]
+                    if image_url is None:
+                        async with session.get(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}",
+                                               headers={"Authorization": f"Bearer {self.config['leonardo-token']}",
+                                                        "Accept": "application/json"},
+                                               ) as response:
+                            response = await response.json()
+                            if "generations_by_pk" not in response:
+                                logging.warning(f"Unexpected response, sleep for {error_sleep}s: {repr(response)}")
+                                await asyncio.sleep(error_sleep)
+                                continue
+                            if len(response["generations_by_pk"]["generated_images"]) == 0:
+                                await asyncio.sleep(error_sleep)
+                                continue
+                            image_url = response["generations_by_pk"]["generated_images"][0]["url"]
+                    if image_bytes is None:
+                        async with session.get(image_url) as response:
+                            image_bytes = BytesIO(await response.read())
+                    async with session.delete(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}",
+                                              headers={"Authorization": f"Bearer {self.config['leonardo-token']}"},
+                                              ) as response:
+                        await response.json()
+                    logging.info(f'Drawed a picture with leonardo AI on this description: {repr(description)}')
+                    return image_bytes
+            except Exception as err:
+                logging.warning(f"Failed to generate image, sleep for {error_sleep}s: {repr(description)}\n{repr(err)}")
+            else:
+                logging.warning(f"Failed to generate image, sleep for {error_sleep}s: {repr(description)}")
+            await asyncio.sleep(error_sleep)
+        raise RuntimeError(f"Failed to generate image {repr(description)}")