From 80e4670e717f2100a94c0d06aa726d16416a5555 Mon Sep 17 00:00:00 2001 From: Hendrik Langer Date: Mon, 3 Apr 2023 12:04:26 +0200 Subject: [PATCH] image generation --- matrix_pygmalion_bot/ai/runpod_pygmalion.py | 6 ++++ matrix_pygmalion_bot/core.py | 32 ++++++++++++++------- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/matrix_pygmalion_bot/ai/runpod_pygmalion.py b/matrix_pygmalion_bot/ai/runpod_pygmalion.py index 30ee8ba..d07bd3f 100644 --- a/matrix_pygmalion_bot/ai/runpod_pygmalion.py +++ b/matrix_pygmalion_bot/ai/runpod_pygmalion.py @@ -219,6 +219,12 @@ async def generate_image(input_prompt: str, negative_prompt: str, api_url: str, return files +async def generate_image1(input_prompt: str, negative_prompt: str, api_key: str): + return await generate_image(input_prompt, negative_prompt, "https://api.runpod.ai/v1/sd-anything-v4/", api_key) + +async def generate_image2(input_prompt: str, negative_prompt: str, api_key: str): + return await generate_image(input_prompt, negative_prompt, "https://api.runpod.ai/v1/sd-openjourney/", api_key) + async def download_image(url, path): r = requests.get(url, stream=True) if r.status_code == 200: diff --git a/matrix_pygmalion_bot/core.py b/matrix_pygmalion_bot/core.py index 5076546..e158feb 100644 --- a/matrix_pygmalion_bot/core.py +++ b/matrix_pygmalion_bot/core.py @@ -82,21 +82,33 @@ class Callbacks(object): print(event) await self.bot.send_message(self.client, room.room_id, "Hello World!") return - elif event.body.startswith('!image ') or event.body.startswith('!image1'): + elif event.body == "!image" or event.body.startswith('!image ') or event.body.startswith('!image1'): prompt = event.body.removeprefix('!image').removeprefix('!image1').strip() negative_prompt = "out of frame, (ugly:1.3), (fused fingers), (too many fingers), (bad anatomy:1.5), (watermark:1.5), (words), letters, untracked eyes, asymmetric eyes, floating head, (logo:1.5), (bad hands:1.3), (mangled hands:1.2), (missing hands), (missing arms), backward hands, floating jewelry, unattached jewelry, floating head, doubled head, unattached head, doubled head, head in body, (misshapen body:1.1), (badly fitted headwear:1.2), floating arms, (too many arms:1.5), limbs fused with body, (facial blemish:1.5), badly fitted clothes, imperfect eyes, untracked eyes, crossed eyes, hair growing from clothes, partial faces, hair not attached to head" - if len(prompt) == 0: - prompt = "a beautiful woman" - output = await ai.generate_image(prompt, negative_prompt, "https://api.runpod.ai/v1/sd-anything-v4/", self.bot.runpod_api_key) + if len(prompt) > 0: + if self.bot.image_prompt: + prompt.replace(self.bot.name, self.bot.image_prompt) + else: + if self.bot.image_prompt: + prompt = self.bot.image_prompt + else: + prompt = "a beautiful woman" + output = await ai.generate_image1(prompt, negative_prompt, self.bot.runpod_api_key) for imagefile in output: await self.bot.send_image(self.client, room.room_id, imagefile) return elif event.body.startswith('!image2'): prompt = event.body.removeprefix('!image2').strip() negative_prompt = "ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face" - if len(prompt) == 0: - prompt = "a beautiful woman" - output = await ai.generate_image(prompt, negative_prompt, "https://api.runpod.ai/v1/sd-openjourney/", self.bot.runpod_api_key) + if len(prompt) > 0: + if self.bot.image_prompt: + prompt.replace(self.bot.name, self.bot.image_prompt) + else: + if self.bot.image_prompt: + prompt = self.bot.image_prompt + else: + prompt = "a beautiful woman" + output = await ai.generate_image2(prompt, negative_prompt, self.bot.runpod_api_key) for imagefile in output: await self.bot.send_image(self.client, room.room_id, imagefile) return @@ -175,9 +187,9 @@ class Callbacks(object): await self.bot.send_message(self.client, room.room_id, translated_answer, reply_to=chat_message.event_id, original_message=answer) else: await self.bot.send_message(self.client, room.room_id, answer, reply_to=chat_message.event_id) - if not "message_count" in self.room_config[room_id]: - self.room_config[room_id]["message_count"] = 0 - self.room_config[room_id]["message_count"] += 1 + if not "message_count" in self.bot.room_config[room.room_id]: + self.bot.room_config[room.room_id]["message_count"] = 0 + self.bot.room_config[room.room_id]["message_count"] += 1 async def invite_cb(self, room: MatrixRoom, event: InviteEvent) -> None: