image4 automatic1111 api

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"<ERROR>")
+
+        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
 
 
 

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