diff --git a/matrix_pygmalion_bot/ai/runpod_pygmalion.py b/matrix_pygmalion_bot/ai/runpod_pygmalion.py index 4f5645e..b81217a 100644 --- a/matrix_pygmalion_bot/ai/runpod_pygmalion.py +++ b/matrix_pygmalion_bot/ai/runpod_pygmalion.py @@ -7,6 +7,10 @@ import requests from transformers import AutoTokenizer, AutoConfig from huggingface_hub import hf_hub_download +import io +import base64 +from PIL import Image, PngImagePlugin + logger = logging.getLogger(__name__) @@ -157,6 +161,13 @@ async def estimate_num_tokens(input_text: str): return len(input_text)//4+1 +async def download_image(url, path): + r = requests.get(url, stream=True) + if r.status_code == 200: + with open(path, 'wb') as f: + for chunk in r: + f.write(chunk) + async def generate_image(input_prompt: str, negative_prompt: str, api_url: str, api_key: str): # Set the API endpoint URL @@ -232,17 +243,74 @@ async def generate_image2(input_prompt: str, negative_prompt: str, api_key: str) async def generate_image3(input_prompt: str, negative_prompt: str, api_key: str): return await generate_image(input_prompt, negative_prompt, "https://api.runpod.ai/v1/mf5f6mocy8bsvx/", api_key) -async def download_image(url, path): - r = requests.get(url, stream=True) - if r.status_code == 200: - with open(path, 'wb') as f: - for chunk in r: - f.write(chunk) +async def generate_image_automatic(input_prompt: str, negative_prompt: str, api_url: str, api_key: str): + + # Set the API endpoint URL + endpoint = api_url + "run" + # Set the headers for the request + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}" + } + # Define your inputs + input_data = { + "input": { + "prompt": input_prompt, + "steps": 5, + }, + } + logger.info(f"sending request to runpod.io") + # Make the request + r = requests.post(endpoint, json=input_data, headers=headers) + r_json = r.json() + logger.info(r_json) + if r.status_code == 200: + status = r_json["status"] + if status != 'IN_QUEUE': + raise ValueError(f"RETURN CODE {status}") + job_id = r_json["id"] + TIMEOUT = 360 + DELAY = 5 + output = None + for i in range(TIMEOUT//DELAY): + endpoint = api_url + "status/" + job_id + r = requests.get(endpoint, headers=headers) + r_json = r.json() + logger.info(r_json) + status = r_json["status"] + if status == 'IN_PROGRESS': + await asyncio.sleep(DELAY) + elif status == 'IN_QUEUE': + await asyncio.sleep(DELAY) + elif status == 'COMPLETED': + output = r_json["output"] + break + else: + raise ValueError(f"RETURN CODE {status}") + + if not output: + raise ValueError(f"") + + os.makedirs("./images", exist_ok=True) + files = [] + for i in output['images']: + temp_name = next(tempfile._get_candidate_names()) + filename = "./images/" + temp_name + ".png" + image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0]))) + info = output['info'] + parameters = output['parameters'] + pnginfo = PngImagePlugin.PngInfo() + pnginfo.add_text("parameters", info) + image.save(filename, pnginfo=pnginfo) + + files.append(filename) + + return files diff --git a/matrix_pygmalion_bot/core.py b/matrix_pygmalion_bot/core.py index 4be049a..4f77513 100644 --- a/matrix_pygmalion_bot/core.py +++ b/matrix_pygmalion_bot/core.py @@ -143,6 +143,26 @@ class Callbacks(object): for imagefile in output: await self.bot.send_image(self.client, room.room_id, imagefile) return + elif event.body.startswith('!image4'): + prompt = event.body.removeprefix('!image4').strip() + negative_prompt = "" + 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" + if self.bot.service == "runpod": + output = await ai.generate_image_automatic(prompt, negative_prompt, "https://api.runpod.ai/v1/lxdhmiccp3vdsf/", self.bot.runpod_api_key) + elif self.bot.service == "stablehorde": + pass + else: + pass + for imagefile in output: + await self.bot.send_image(self.client, room.room_id, imagefile) + return elif event.body.startswith('!begin'): self.bot.chat_history.room(room.display_name).clear() self.bot.room_config[room.room_id]["tick"] = 0