Fjerkroa/discord_bot
1
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Add possibility to use Leonardo.ai for image generation and improve it a bit.

Browse Source
This commit is contained in:
OK
2023-08-19 16:04:21 +02:00
parent e237800348
commit 4ec20f1b26
1 changed files with 44 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fjerkroa_bot/ai_responder.py
+44 -9
View File
@@ -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)