From 4ec20f1b264055f06a0f7808b4380f51e57539fc Mon Sep 17 00:00:00 2001
From: Oleksandr Kozachuk <ddeus.lp@mailnull.com>
Date: Sat, 19 Aug 2023 16:04:21 +0200
Subject: [PATCH] Add possibility to use Leonardo.ai for image generation and
 improve it a bit.

---
 fjerkroa_bot/ai_responder.py | 53 ++++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 9 deletions(-)

diff --git a/fjerkroa_bot/ai_responder.py b/fjerkroa_bot/ai_responder.py
index cafdf55..13dcb47 100644
--- a/fjerkroa_bot/ai_responder.py
+++ b/fjerkroa_bot/ai_responder.py
@@ -163,23 +163,37 @@ class AIResponder(object):
                                             json={"prompt": description,
                                                   "modelId": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3",
                                                   "num_images": 1,
+                                                  "sd_version": "v2",
+                                                  "promptMagic": True,
                                                   "width": 512,
                                                   "height": 512},
                                             headers={"Authorization": f"Bearer {self.config['leonardo-token']}",
                                                      "Accept": "application/json",
                                                      "Content-Type": "application/json"},
                                             ) as response:
-                        generation_id = (await response.json())["sdGenerationJob"]["generationId"]
-                    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:
-                        image_url = (await response.json())["generations_by_pk"]["generated_images"][0]["url"]
+                        response = await response.json()
+                        generation_id = response["sdGenerationJob"]["generationId"]
+                    while True:
+                        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 len(response["generations_by_pk"]["generated_images"]) == 0:
+                                await asyncio.sleep(0.1)
+                                continue
+                            image_url = response["generations_by_pk"]["generated_images"][0]["url"]
+                            break
                     async with session.get(image_url) as response:
-                        return BytesIO(await response.read())
+                        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()
+                    return image_bytes
             except Exception as err:
-                logging.warning(f"Failed to generate image {repr(description)}: {repr(err)}")
-        raise RuntimeError(f"Failed to generate image {repr(description)} after multiple retries")
+                logging.warning(f"Failed to generate image: {repr(err)}")
+        raise RuntimeError(f"Failed to generate image {repr(description)}")
 
     async def post_process(self, message: AIMessage, response: Dict[str, Any]) -> AIResponse:
         for fld in ('answer', 'channel', 'staff', 'picture', 'hack'):
@@ -276,6 +290,24 @@ class AIResponder(object):
             logging.warning(f"failed to execute a fix for the answer: {repr(err)}")
         return answer
 
+    async def translate(self, text: str, language: str = "english") -> str:
+        if 'fix-model' not in self.config:
+            return text
+        message = [{"role": "system", "content": f"You are an professional translator to {language} language,"
+                                                 f" you translate everything you get directly to {language}."},
+                   {"role": "user", "content": text}]
+        try:
+            result = await openai.ChatCompletion.acreate(model=self.config["fix-model"],
+                                                         messages=message,
+                                                         temperature=0.2,
+                                                         max_tokens=2048)
+            response = result['choices'][0]['message']['content']
+            logging.info(f"got this translated message:\n{pp(response)}")
+            return response
+        except Exception as err:
+            logging.warning(f"failed to translate the text: {repr(err)}")
+        return text
+
     def shrink_history_by_one(self, index: int = 0) -> None:
         if index >= len(self.history):
             del self.history[0]
@@ -339,6 +371,9 @@ class AIResponder(object):
                 retries -= 1
                 continue
 
+            if response.get("picture") is not None:
+                response["picture"] = await self.translate(response["picture"])
+
             # Post-process the message and update the answer's content
             answer_message = await self.post_process(message, response)
             answer['content'] = str(answer_message)