stablehorde

matrix_pygmalion_bot/ai/runpod_pygmalion.py

@@ -174,7 +174,7 @@ async def generate_image(input_prompt: str, negative_prompt: str, api_url: str,
             "prompt": input_prompt,
             "negative_prompt": negative_prompt,
             "width": 512,
-            "height": 768,
+            "height": 512,
             "nsfw": True
         }
     }
@@ -193,6 +193,7 @@ async def generate_image(input_prompt: str, negative_prompt: str, api_url: str,
         job_id = r_json["id"]
         TIMEOUT = 180
         DELAY = 5
+        output = None
         for i in range(TIMEOUT//DELAY):
             endpoint = api_url + "status/" + job_id
             r = requests.get(endpoint, headers=headers)
@@ -209,6 +210,9 @@ async def generate_image(input_prompt: str, negative_prompt: str, api_url: str,
             else:
                 raise ValueError(f"RETURN CODE {status}")
 
+        if not output:
+            raise ValueError(f"<ERROR>")
+
         os.makedirs("./images", exist_ok=True)
         files = []
         for image in output:
@@ -225,6 +229,9 @@ async def generate_image1(input_prompt: str, negative_prompt: str, api_key: str)
 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 generate_image3(input_prompt: str, negative_prompt: str, api_key: str):
+    return await generate_image(input_prompt, negative_prompt, "https://api.runpod.ai/v1/r9em5b11dkzvo0/", api_key)
+
 async def download_image(url, path):
     r = requests.get(url, stream=True)
     if r.status_code == 200:

matrix_pygmalion_bot/ai/stablehorde.py

@@ -0,0 +1,256 @@
+import asyncio
+import os, tempfile
+import logging
+
+import requests
+
+from transformers import AutoTokenizer, AutoConfig
+from huggingface_hub import hf_hub_download
+
+logger = logging.getLogger(__name__)
+
+
+async def generate_sync(
+    prompt: str,
+    api_key: str,
+    bot_name: str,
+):
+    # Set the API endpoint URL
+    endpoint = "https://koboldai.net/api/v2/generate/async"
+
+    # Set the headers for the request
+    headers = {
+        "Content-Type": "application/json",
+        "accept": "application/json",
+        "apikey": f"{api_key}"
+    }
+
+    max_new_tokens = 200
+    prompt_num_tokens = await num_tokens(prompt)
+
+    # Define your inputs
+    input_data = {
+        "prompt": prompt,
+        "params": {
+            "n": 1,
+#            "frmtadsnsp": False,
+#            "frmtrmblln": False,
+#            "frmtrmspch": False,
+#            "frmttriminc": False,
+            "max_context_length": 1024,
+            "max_length": 512,
+            "rep_pen": 1.1,
+            "rep_pen_range": 1024,
+            "rep_pen_slope": 0.7,
+#            "singleline": False,
+#            "soft_prompt": "",
+            "temperature": 0.75,
+            "tfs": 1.0,
+            "top_a": 0.0,
+            "top_k": 0,
+            "top_p": 0.9,
+            "typical": 1.0,
+#            "sampler_order": [0],
+        },
+        "softprompts": [],
+        "trusted_workers": False,
+        "nsfw": True,
+#        "workers": [],
+        "models": ["PygmalionAI/pygmalion-6b"]
+    }
+
+
+    logger.info(f"sending request to koboldai.net")
+
+    # Make the request
+    r = requests.post(endpoint, json=input_data, headers=headers, timeout=180)
+
+    r_json = r.json()
+    logger.info(r_json)
+    status = r_json["message"]
+
+    if "id" in r_json:
+        job_id = r_json["id"]
+        TIMEOUT = 360
+        DELAY = 11
+        for i in range(TIMEOUT//DELAY):
+            endpoint = "https://koboldai.net/api/v2/generate/text/status/" + job_id
+            r = requests.get(endpoint, headers=headers)
+            r_json = r.json()
+            logger.info(r_json)
+            if "done" not in r_json:
+                return "<ERROR>"
+            if r_json["done"] == True:
+                text = r_json["generations"][0]["text"]
+                answer = text.removeprefix(prompt)
+                idx = answer.find(f"\nYou:")
+                if idx != -1:
+                    reply = answer[:idx].strip()
+                else:
+                    reply = answer.removesuffix('<|endoftext|>').strip()
+                reply = reply.replace("\n{bot_name}: ", " ")
+                reply.replace("\n<BOT>: ", " ")
+                return reply
+            else:
+                await asyncio.sleep(DELAY)
+    else:
+        return "<ERROR> {status}"
+
+async def get_full_prompt(simple_prompt: str, bot, chat_history):
+
+    # Prompt without history
+    prompt = bot.name + "'s Persona: " + bot.persona + "\n"
+    prompt += "Scenario: " + bot.scenario + "\n"
+    prompt += "<START>" + "\n"
+    #prompt += bot.name + ": " + bot.greeting + "\n"
+    prompt += "You: " + simple_prompt + "\n"
+    prompt += bot.name + ":"
+
+    MAX_TOKENS = 2048
+    max_new_tokens = 200
+    total_num_tokens = await num_tokens(prompt)
+    visible_history = []
+    current_message = True
+    for key, chat_item in reversed(chat_history.chat_history.items()):
+        if current_message:
+            current_message = False
+            continue
+        if chat_item.message["en"].startswith('!begin'):
+            break
+        if chat_item.message["en"].startswith('!'):
+            continue
+        if chat_item.message["en"].startswith('<ERROR>'):
+            continue
+        #if chat_item.message["en"] == bot.greeting:
+        #    continue
+        if chat_item.num_tokens == None:
+            chat_item.num_tokens = await num_tokens("{}: {}".format(chat_item.user_name, chat_item.message["en"]))
+        # TODO: is it MAX_TOKENS or MAX_TOKENS - max_new_tokens??
+        logger.debug(f"History: " + str(chat_item) + " [" + str(chat_item.num_tokens) + "]")
+        if total_num_tokens + chat_item.num_tokens < MAX_TOKENS - max_new_tokens:
+            visible_history.append(chat_item)
+            total_num_tokens += chat_item.num_tokens
+        else:
+            break
+    visible_history = reversed(visible_history)
+
+    prompt = bot.name + "'s Persona: " + bot.persona + "\n"
+    prompt += "Scenario: " + bot.scenario + "\n"
+    prompt += "<START>" + "\n"
+    #prompt += bot.name + ": " + bot.greeting + "\n"
+    for chat_item in visible_history:
+        if chat_item.is_own_message:
+            prompt += bot.name + ": " + chat_item.message["en"] + "\n"
+        else:
+            prompt += "You" + ": " + chat_item.message["en"] + "\n"
+    prompt += "You: " + simple_prompt + "\n"
+    prompt += bot.name + ":"
+
+    return prompt
+
+
+async def num_tokens(input_text: str):
+#    os.makedirs("./models/pygmalion-6b", exist_ok=True)
+#    hf_hub_download(repo_id="PygmalionAI/pygmalion-6b", filename="config.json", cache_dir="./models/pygmalion-6b")
+#    config = AutoConfig.from_pretrained("./models/pygmalion-6b/config.json")
+    tokenizer = AutoTokenizer.from_pretrained("PygmalionAI/pygmalion-6b")
+    encoding = tokenizer.encode(input_text, add_special_tokens=False)
+    max_input_size = tokenizer.max_model_input_sizes
+    return len(encoding)
+
+async def estimate_num_tokens(input_text: str):
+    return len(input_text)//4+1
+
+
+async def generate_image(input_prompt: str, negative_prompt: str, model: str, api_key: str):
+
+    # Set the API endpoint URL
+    endpoint = "https://stablehorde.net/api/v2/generate/async"
+
+    # Set the headers for the request
+    headers = {
+        "Content-Type": "application/json",
+        "accept": "application/json",
+        "apikey": f"{api_key}"
+    }
+
+    # Define your inputs
+    input_data = {
+        "prompt": input_prompt,
+        "params": {
+#            "negative_prompt": negative_prompt,
+            "width": 512,
+            "height": 512,
+        },
+        "nsfw": True,
+        "trusted_workers": False,
+#        "workers": [],
+        "models": ["{model}"]
+    }
+
+    logger.info(f"sending request to stablehorde.net")
+
+    # Make the request
+    r = requests.post(endpoint, json=input_data, headers=headers)
+    r_json = r.json()
+    logger.info(r_json)
+
+    if r.status_code == 202:
+        #status = r_json["message"]
+        job_id = r_json["id"]
+        TIMEOUT = 360
+        DELAY = 11
+        output = None
+        for i in range(TIMEOUT//DELAY):
+            endpoint = "https://stablehorde.net/api/v2/generate/status/" + job_id
+            r = requests.get(endpoint, headers=headers)
+            r_json = r.json()
+            logger.info(r_json)
+            #status = r_json["message"]
+            if "done" not in r_json:
+                return "<ERROR>"
+            if "faulted" in r_json and r_json["faulted"] == True:
+                return "<ERROR>"
+            if r_json["done"] == True:
+                output = r_json["generations"]
+                break
+            else:
+                await asyncio.sleep(DELAY)
+
+        if not output:
+            raise ValueError(f"<ERROR>")
+
+        os.makedirs("./images", exist_ok=True)
+        files = []
+        for image in output:
+            temp_name = next(tempfile._get_candidate_names())
+            filename = "./images/" + temp_name + ".jpg"
+            await download_image(image["img"], filename)
+            files.append(filename)
+
+        return files
+
+async def generate_image1(input_prompt: str, negative_prompt: str, api_key: str):
+    return await generate_image(input_prompt, negative_prompt, "Deliberate", api_key)
+
+async def generate_image2(input_prompt: str, negative_prompt: str, api_key: str):
+    return await generate_image(input_prompt, negative_prompt, "PFG", api_key)
+
+async def generate_image3(input_prompt: str, negative_prompt: str, api_key: str):
+    return await generate_image(input_prompt, negative_prompt, "Hassanblend", 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)
+
+
+
+
+
+
+
+
+

matrix_pygmalion_bot/core.py

@@ -16,7 +16,8 @@ import json
 
 from .helpers import Event
 from .chatlog import BotChatHistory
-ai = importlib.import_module("matrix_pygmalion_bot.ai.runpod_pygmalion")
+#ai = importlib.import_module("matrix_pygmalion_bot.ai.runpod_pygmalion")
+ai = importlib.import_module("matrix_pygmalion_bot.ai.stablehorde")
 #from .llama_cpp import generate, get_full_prompt, get_full_prompt_chat_style
 #from .runpod_pygmalion import generate_sync, get_full_prompt
 import matrix_pygmalion_bot.translate as translate
@@ -93,7 +94,12 @@ class Callbacks(object):
                     prompt = self.bot.image_prompt
                 else:
                     prompt = "a beautiful woman"
-            output = await ai.generate_image1(prompt, negative_prompt, self.bot.runpod_api_key)
+            if self.bot.service == "runpod":
+                output = await ai.generate_image1(prompt, negative_prompt, self.bot.runpod_api_key)
+            elif self.bot.service == "stablehorde":
+                output = await ai.generate_image1(prompt, negative_prompt, self.bot.stablehorde_api_key)
+            else:
+                pass
             for imagefile in output:
                 await self.bot.send_image(self.client, room.room_id, imagefile)
             return
@@ -108,7 +114,32 @@ class Callbacks(object):
                     prompt = self.bot.image_prompt
                 else:
                     prompt = "a beautiful woman"
-            output = await ai.generate_image2(prompt, negative_prompt, self.bot.runpod_api_key)
+            if self.bot.service == "runpod":
+                output = await ai.generate_image2(prompt, negative_prompt, self.bot.runpod_api_key)
+            elif self.bot.service == "stablehorde":
+                output = await ai.generate_image2(prompt, negative_prompt, self.bot.stablehorde_api_key)
+            else:
+                pass
+            for imagefile in output:
+                await self.bot.send_image(self.client, room.room_id, imagefile)
+            return
+        elif event.body.startswith('!image3'):
+            prompt = event.body.removeprefix('!image3').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:
+                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_image3(prompt, negative_prompt, self.bot.runpod_api_key)
+            elif self.bot.service == "stablehorde":
+                output = await ai.generate_image3(prompt, negative_prompt, self.bot.stablehorde_api_key)
+            else:
+                pass
             for imagefile in output:
                 await self.bot.send_image(self.client, room.room_id, imagefile)
             return
@@ -385,8 +416,12 @@ async def main() -> None:
             events = config[section]['events'].strip().split('\n')
             for event in events:
                 await bot.add_event(event)
+        if config.has_option('DEFAULT', 'service'):
+            bot.service = config['DEFAULT']['service']
         if config.has_option('DEFAULT', 'runpod_api_key'):
             bot.runpod_api_key = config['DEFAULT']['runpod_api_key']
+        if config.has_option('DEFAULT', 'stablehorde_api_key'):
+            bot.stablehorde_api_key = config['DEFAULT']['stablehorde_api_key']
         await bot.read_conf2(section)
         bots.append(bot)
         await bot.login()